Activate ANN001 on tests/integrations (#5590)

Author: cbeauchesne

pyproject.toml

@@ -174,8 +174,7 @@ ignore = [
 "tests/*test_*.py" = [
     "ANN201",    # missing-return-type-undocumented-public-function: allow test_method to not declare their return type
 ]
-
-"tests/{parametric,appsec,debugger,remote_config,integrations}*/test_*.py" = [
+"tests/{parametric,appsec,debugger,remote_config}*/test_*.py" = [
     "ANN001",    # missing-type-function-argument: TODO
 ]
 "tests/*" = [

tests/integrations/crossed_integrations/test_kafka.py

@@ -15,7 +15,7 @@ class _BaseKafka:
     buddy_interface: interfaces.LibraryInterfaceValidator
 
     @classmethod
-    def get_span(cls, interface, span_kind, topic) -> dict | None:
+    def get_span(cls, interface: interfaces.LibraryInterfaceValidator, span_kind: str, topic: str) -> dict | None:
         logger.debug(f"Trying to find traces with span kind: {span_kind} and topic: {topic} in {interface}")
 
         for data, trace in interface.get_traces():
@@ -36,7 +36,7 @@ def get_span(cls, interface, span_kind, topic) -> dict | None:
         return None
 
     @staticmethod
-    def get_topic(span) -> str | None:
+    def get_topic(span: dict) -> str | None:
         """Extracts the topic from a span by trying various fields"""
         topic = span["meta"].get("kafka.topic")  # this is in python
         if topic is None:
@@ -131,7 +131,12 @@ def test_consume_trace_equality(self):
         assert consumer_span is not None
         assert producer_span["trace_id"] == consumer_span["trace_id"]
 
-    def validate_kafka_spans(self, producer_interface, consumer_interface, topic):
+    def validate_kafka_spans(
+        self,
+        producer_interface: interfaces.LibraryInterfaceValidator,
+        consumer_interface: interfaces.LibraryInterfaceValidator,
+        topic: str,
+    ):
         """Validates production/consumption of kafka message.
         It works the same for both test_produce and test_consume
         """

tests/integrations/crossed_integrations/test_kinesis.py

@@ -15,7 +15,9 @@ class _BaseKinesis:
     unique_id: str
 
     @classmethod
-    def get_span(cls, interface, span_kind, stream, operation) -> dict | None:
+    def get_span(
+        cls, interface: interfaces.LibraryInterfaceValidator, span_kind: list[str], stream: str, operation: str
+    ) -> dict | None:
         logger.debug(f"Trying to find traces with span kind: {span_kind} and stream: {stream} in {interface}")
 
         for data, trace in interface.get_traces():
@@ -53,7 +55,7 @@ def get_span(cls, interface, span_kind, stream, operation) -> dict | None:
         return None
 
     @staticmethod
-    def get_stream(span) -> str | None:
+    def get_stream(span: dict) -> str | None:
         """Extracts the stream from a span by trying various fields"""
         stream = span["meta"].get("streamname", None)  # this is in nodejs, java, python
 
@@ -187,7 +189,12 @@ def test_consume_trace_equality(self):
         assert consumer_span is not None
         assert producer_span["trace_id"] == consumer_span["trace_id"]
 
-    def validate_kinesis_spans(self, producer_interface, consumer_interface, stream):
+    def validate_kinesis_spans(
+        self,
+        producer_interface: interfaces.LibraryInterfaceValidator,
+        consumer_interface: interfaces.LibraryInterfaceValidator,
+        stream: str,
+    ):
         """Validates production/consumption of Kinesis message.
         It works the same for both test_produce and test_consume
         """

tests/integrations/crossed_integrations/test_rabbitmq.py

@@ -19,7 +19,14 @@ class _BaseRabbitMQ:
     buddy_interface: interfaces.LibraryInterfaceValidator
 
     @classmethod
-    def get_span(cls, interface, span_kind, queue, exchange, operation) -> dict | None:
+    def get_span(
+        cls,
+        interface: interfaces.LibraryInterfaceValidator,
+        span_kind: str,
+        queue: str,
+        exchange: str,
+        operation: list[str],
+    ) -> dict | None:
         logger.debug(f"Trying to find traces with span kind: {span_kind} and queue: {queue} in {interface}")
 
         for data, trace in interface.get_traces():
@@ -32,17 +39,17 @@ def get_span(cls, interface, span_kind, queue, exchange, operation) -> dict | No
 
                 operation_found = False
                 for op in operation:
-                    if op.lower() in span.get("resource").lower() or op.lower() in span.get("name").lower():
+                    if op.lower() in span["resource"].lower() or op.lower() in span["name"].lower():
                         operation_found = True
                         break
 
                 if not operation_found:
                     continue
 
-                meta = span.get("meta")
+                meta: dict[str, str] = span["meta"]
                 if (
-                    queue.lower() not in span.get("resource").lower()
-                    and exchange.lower() not in span.get("resource").lower()
+                    queue.lower() not in span["resource"].lower()
+                    and exchange.lower() not in span["resource"].lower()
                     and queue.lower() not in meta.get("rabbitmq.routing_key", "").lower()
                     # this is where we find the queue name in dotnet 👇
                     and queue.lower() not in meta.get("amqp.routing_key", "").lower()
@@ -188,7 +195,13 @@ def test_consume_trace_equality(self):
         assert consumer_span is not None
         assert producer_span["trace_id"] == consumer_span["trace_id"]
 
-    def validate_rabbitmq_spans(self, producer_interface, consumer_interface, queue, exchange):
+    def validate_rabbitmq_spans(
+        self,
+        producer_interface: interfaces.LibraryInterfaceValidator,
+        consumer_interface: interfaces.LibraryInterfaceValidator,
+        queue: str,
+        exchange: str,
+    ):
         """Validates production/consumption of RabbitMQ message.
         It works the same for both test_produce and test_consume
         """

tests/integrations/crossed_integrations/test_sns_to_sqs.py

@@ -17,7 +17,14 @@ class _BaseSNS:
     unique_id: str
 
     @classmethod
-    def get_span(cls, interface, span_kind, queue, topic, operation) -> dict | None:
+    def get_span(
+        cls,
+        interface: interfaces.LibraryInterfaceValidator,
+        span_kind: list[str],
+        queue: str,
+        topic: str,
+        operation: str,
+    ) -> dict | None:
         logger.debug(f"Trying to find traces with span kind: {span_kind} and queue: {queue} in {interface}")
         manual_span_found = False
 
@@ -67,7 +74,7 @@ def get_span(cls, interface, span_kind, queue, topic, operation) -> dict | None:
         return None
 
     @staticmethod
-    def get_queue(span) -> str | None:
+    def get_queue(span: dict) -> str | None:
         """Extracts the queue from a span by trying various fields"""
         queue = span["meta"].get("queuename", None)  # this is in nodejs, java, python
 
@@ -83,7 +90,7 @@ def get_queue(span) -> str | None:
         return queue
 
     @staticmethod
-    def get_topic(span) -> str | None:
+    def get_topic(span: dict) -> str | None:
         """Extracts the topic from a span by trying various fields"""
         topic = span["meta"].get("topicname", None)  # this is in nodejs, java, python
 
@@ -217,7 +224,13 @@ def test_consume_trace_equality(self):
         assert consumer_span is not None
         assert producer_span["trace_id"] == consumer_span["trace_id"]
 
-    def validate_sns_spans(self, producer_interface, consumer_interface, queue, topic):
+    def validate_sns_spans(
+        self,
+        producer_interface: interfaces.LibraryInterfaceValidator,
+        consumer_interface: interfaces.LibraryInterfaceValidator,
+        queue: str,
+        topic: str,
+    ):
         """Validates production/consumption of sns message.
         It works the same for both test_produce and test_consume
         """

tests/integrations/crossed_integrations/test_sqs.py

@@ -15,7 +15,9 @@ class _BaseSQS:
     unique_id: str
 
     @classmethod
-    def get_span(cls, interface, span_kind, queue, operation) -> dict | None:
+    def get_span(
+        cls, interface: interfaces.LibraryInterfaceValidator, span_kind: list[str], queue: str, operation: str
+    ) -> dict | None:
         logger.debug(f"Trying to find traces with span kind: {span_kind} and queue: {queue} in {interface}")
         manual_span_found = False
 
@@ -66,7 +68,7 @@ def get_span(cls, interface, span_kind, queue, operation) -> dict | None:
         return None
 
     @staticmethod
-    def get_queue(span) -> str | None:
+    def get_queue(span: dict) -> str | None:
         """Extracts the queue from a span by trying various fields"""
         queue = span["meta"].get("queuename", None)  # this is in nodejs, java, python
 
@@ -200,7 +202,12 @@ def test_consume_trace_equality(self):
         assert consumer_span is not None
         assert producer_span["trace_id"] == consumer_span["trace_id"]
 
-    def validate_sqs_spans(self, producer_interface, consumer_interface, queue):
+    def validate_sqs_spans(
+        self,
+        producer_interface: interfaces.LibraryInterfaceValidator,
+        consumer_interface: interfaces.LibraryInterfaceValidator,
+        queue: str,
+    ):
         """Validates production/consumption of sqs message.
         It works the same for both test_produce and test_consume
         """

tests/integrations/test_db_integrations_sql.py

@@ -12,7 +12,7 @@ class _BaseDatadogDbIntegrationTestClass(BaseDbIntegrationsTestClass):
     Check integration spans status: https://docs.google.com/spreadsheets/d/1qm3B0tJ-gG11j_MHoEd9iMXf4_DvWAGCLwmBhWCxbA8/edit#gid=623219645
     """
 
-    def get_spans(self, excluded_operations=(), operations=None):
+    def get_spans(self, excluded_operations: tuple[str, ...] = (), operations: list[str] | None = None):
         """Get the spans from tracer and agent generated by all requests"""
 
         # yield the span from the tracer in first, as if it fails, there is a good chance that the one from the agent also fails
@@ -26,18 +26,18 @@ def get_spans(self, excluded_operations=(), operations=None):
             yield db_operation, self.get_span_from_agent(request)
 
     # Tests methods
-    def test_sql_traces(self, excluded_operations=()):
+    def test_sql_traces(self, excluded_operations: tuple[str, ...] = ()):
         """After make the requests we check that we are producing sql traces"""
         for _, span in self.get_spans(excluded_operations):
             assert span is not None
 
     def test_resource(self):
         """Usually the query"""
 
-        for db_operation, span in self.get_spans(excluded_operations=["procedure", "select_error"]):
+        for db_operation, span in self.get_spans(excluded_operations=("procedure", "select_error")):
             assert db_operation in span["resource"].lower()
 
-    def test_sql_success(self, excluded_operations=()):
+    def test_sql_success(self, excluded_operations: tuple[str, ...] = ()):
         """We check all sql launched for the app work"""
 
         for db_operation, span in self.get_spans(excluded_operations=excluded_operations + ("select_error",)):
@@ -49,7 +49,7 @@ def test_sql_success(self, excluded_operations=()):
                 raise ValueError(f"Error found in {db_operation} operation, please check captured log call")
 
     @irrelevant(library="python", reason="Python is using the correct span: db.system")
-    def test_db_type(self, excluded_operations=()):
+    def test_db_type(self, excluded_operations: tuple[str, ...] = ()):
         """DEPRECATED!! Now it is db.system. An identifier for the database management system (DBMS) product being used.
         Must be one of the available values: https://datadoghq.atlassian.net/wiki/spaces/APM/pages/2357395856/Span+attributes#db.system
         """
@@ -65,7 +65,7 @@ def test_db_name(self):
         for db_operation, span in self.get_spans():
             assert span["meta"]["db.name"] == db_container.db_instance, f"Test is failing for {db_operation}"
 
-    def test_span_kind(self, excluded_operations=()):
+    def test_span_kind(self, excluded_operations: tuple[str, ...] = ()):
         """Describes the relationship between the Span, its parents, and its children in a Trace."""
 
         for _, span in self.get_spans(excluded_operations):
@@ -98,7 +98,7 @@ def test_db_connection_string(self):
         for db_operation, span in self.get_spans():
             assert span["meta"]["db.connection_string"].strip(), f"Test is failing for {db_operation}"
 
-    def test_db_user(self, excluded_operations=()):
+    def test_db_user(self, excluded_operations: tuple[str, ...] = ()):
         """Username for accessing the database."""
         db_container = context.get_container_by_dd_integration_name(self.db_service)
 
@@ -107,7 +107,7 @@ def test_db_user(self, excluded_operations=()):
 
     @missing_feature(library="python", reason="not implemented yet")
     @missing_feature(library="nodejs", reason="not implemented yet")
-    def test_db_instance(self, excluded_operations=()):
+    def test_db_instance(self, excluded_operations: tuple[str, ...] = ()):
         """The name of the database being connected to. Database instance name. Formerly db.name"""
         db_container = context.get_container_by_dd_integration_name(self.db_service)
 
@@ -120,12 +120,12 @@ def test_db_instance(self, excluded_operations=()):
     def test_db_statement_query(self):
         """Usually the query"""
 
-        for db_operation, span in self.get_spans(excluded_operations=["procedure", "select_error"]):
+        for db_operation, span in self.get_spans(excluded_operations=("procedure", "select_error")):
             assert db_operation in span["meta"]["db.statement"].lower()
 
     @missing_feature(library="nodejs", reason="not implemented yet")
     @missing_feature(library="python", reason="not implemented yet")
-    def test_db_operation(self, excluded_operations=()):
+    def test_db_operation(self, excluded_operations: tuple[str, ...] = ()):
         """The name of the operation being executed"""
 
         for db_operation, span in self.get_spans(excluded_operations=excluded_operations + ("select_error",)):
@@ -144,7 +144,7 @@ def test_db_operation(self, excluded_operations=()):
     def test_db_sql_table(self):
         """The name of the primary table that the operation is acting upon, including the database name (if applicable)."""
 
-        for db_operation, span in self.get_spans(excluded_operations=["procedure"]):
+        for db_operation, span in self.get_spans(excluded_operations=("procedure",)):
             assert span["meta"]["db.sql.table"].strip(), f"Test is failing for {db_operation}"
 
     @missing_feature(library="python", reason="not implemented yet")
@@ -158,7 +158,7 @@ def test_db_row_count(self):
         for _, span in self.get_spans(operations=["select"]):
             assert span["meta"]["db.row_count"] > 0, "Test is failing for select"
 
-    def test_db_password(self, excluded_operations=()):
+    def test_db_password(self, excluded_operations: tuple[str, ...] = ()):
         """The database password should not show in the traces"""
         db_container = context.get_container_by_dd_integration_name(self.db_service)
 
@@ -214,7 +214,7 @@ def test_sql_query(self):
                 db_operation in span["meta"]["sql.query"].lower()
             ), f"sql.query span not found for operation {db_operation}"
 
-    def test_obfuscate_query(self, excluded_operations=()):
+    def test_obfuscate_query(self, excluded_operations: tuple[str, ...] = ()):
         """All queries come out obfuscated from agent"""
         for db_operation, request in self.get_requests(excluded_operations=excluded_operations):
             span = self.get_span_from_agent(request)
@@ -295,7 +295,7 @@ def test_db_system(self):
     def test_db_user(self):
         super().test_db_user()
 
-    def test_obfuscate_query(self, excluded_operations=()):
+    def test_obfuscate_query(self, excluded_operations: tuple[str, ...] = ()):
         """All queries come out obfuscated from agent"""
         for db_operation, request in self.get_requests(excluded_operations=excluded_operations):
             span = self.get_span_from_agent(request)

tests/integrations/test_dbm.py

@@ -7,9 +7,10 @@
 import re
 
 from utils import weblog, interfaces, context, scenarios, features, irrelevant, flaky, bug, logger
+from utils._weblog import HttpResponse
 
 
-def remove_traceparent(s):
+def remove_traceparent(s: str) -> str:
     return re.sub(r",traceparent='[^']*'", "", s)
 
 
@@ -53,7 +54,7 @@ def weblog_trace_payload(self):
                     ]
                 )
 
-    def _get_db_span(self, response):
+    def _get_db_span(self, response: HttpResponse) -> dict:
         assert response.status_code == 200, f"Request: {context.scenario.name} wasn't successful."
 
         spans = []
@@ -82,11 +83,11 @@ def _assert_spans_are_untagged(self):
         for request in self.requests:
             self._assert_span_is_untagged(self._get_db_span(request))
 
-    def _assert_span_is_untagged(self, span):
+    def _assert_span_is_untagged(self, span: dict) -> None:
         meta = span.get("meta", {})
         assert self.META_TAG not in meta, f"{self.META_TAG} found in span meta: {json.dumps(span, indent=2)}"
 
-    def _assert_span_is_tagged(self, span):
+    def _assert_span_is_tagged(self, span: dict) -> None:
         meta = span.get("meta", {})
         assert self.META_TAG in meta, f"{self.META_TAG} not found in span meta: {json.dumps(span, indent=2)}"
         tag_value = meta.get(self.META_TAG)

tests/integrations/test_dsm.py

@@ -56,7 +56,7 @@
 WEBLOG_VARIANT_SANITIZED = context.weblog_variant.replace(".", "_").replace(" ", "_").replace("/", "_")
 
 
-def get_message(test, system):
+def get_message(test: str, system: str) -> str:
     return f"[test_dsm.py::{test}] [{system.upper()}] Hello from {context.library.name} DSM test: {scenarios.integrations_aws.unique_id}"
 
 
@@ -636,14 +636,14 @@ def test_dsm_manual_checkpoint_inter_process(self):
 
 class DsmHelper:
     @staticmethod
-    def is_tags_included(actual_tags, expected_tags) -> bool:
+    def is_tags_included(actual_tags: tuple, expected_tags: tuple) -> bool:
         assert isinstance(actual_tags, tuple)
         assert isinstance(expected_tags, tuple)
 
         return all(expected_tag in actual_tags for expected_tag in expected_tags)
 
     @staticmethod
-    def assert_checkpoint_presence(hash_, parent_hash, tags) -> None:
+    def assert_checkpoint_presence(hash_: int, parent_hash: int, tags: tuple) -> None:
         assert isinstance(tags, tuple)
 
         logger.info(f"Look for {hash_}, {parent_hash}, {tags}")