Add more metrics, consumer, and producer tests (#619)

* Add more metrics, consumer, and producer tests

Co-authored-by: Timothy Pansino <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>

* Remove breakpoint

* Add DT accepted validator

* Fix lint error

* Fix issue with holdover transaction

Co-authored-by: Lalleh Rafeei <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>

* [Mega-Linter] Apply linters fixes

* Fix broken producer error test

* [Mega-Linter] Apply linters fixes

* Bump Tests

* Fix flakey test

* Fix flakey test

* fixup

* Fix metrics tests for Python 2.7

Co-authored-by: Timothy Pansino <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: Uma Annamalai <[email protected]>

Co-authored-by: Timothy Pansino <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: Tim Pansino <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>
Co-authored-by: Uma Annamalai <[email protected]>

Author: 4 people

.github/workflows/tests.yml

@@ -447,15 +447,15 @@ jobs:
 
     services:
       zookeeper:
-        image: bitnami/zookeeper:latest
+        image: bitnami/zookeeper:3.7
         env:
           ALLOW_ANONYMOUS_LOGIN: yes
 
         ports:
           - 2181:2181
 
       kafka:
-        image: bitnami/kafka:latest
+        image: bitnami/kafka:3.2
         ports:
           - 8080:8080
           - 8081:8081

newrelic/hooks/messagebroker_kafkapython.py

@@ -158,7 +158,7 @@ def wrap_kafkaconsumer_next(wrapped, instance, args, kwargs):
                 library=library,
                 destination_type=destination_type,
                 destination_name=destination_name,
-                headers=record.headers,
+                headers=dict(record.headers),
                 transport_type="Kafka",
                 routing_key=record.key,
                 source=wrapped,

tests/messagebroker_kafkapython/conftest.py

@@ -55,7 +55,7 @@
 
 
 @pytest.fixture(scope="function")
-def producer(data_source):
+def producer(topic, data_source):
     producer = kafka.KafkaProducer(
         bootstrap_servers=BROKER, api_version=(2, 0, 2), value_serializer=lambda v: json.dumps(v).encode("utf-8")
     )
@@ -64,23 +64,39 @@ def producer(data_source):
 
 
 @pytest.fixture(scope="function")
-def consumer(topic, data_source):
+def consumer(topic, data_source, producer):
     consumer = kafka.KafkaConsumer(
         topic,
         bootstrap_servers=BROKER,
         value_deserializer=lambda v: json.loads(v.decode("utf-8")),
         auto_offset_reset="earliest",
-        consumer_timeout_ms=5000,
+        consumer_timeout_ms=500,
         heartbeat_interval_ms=1000,
         group_id="test",
     )
+    # The first time the kafka consumer is created and polled, it returns a StopIterator
+    # exception. To by-pass this, loop over the consumer before using it.
+    # NOTE: This seems to only happen in Python2.7.
+    for record in consumer:
+        pass
     yield consumer
     consumer.close()
 
 
 @pytest.fixture(scope="function")
 def topic():
-    yield "test-topic-%s" % str(uuid.uuid4())
+    # from kafka.admin.client import KafkaAdminClient
+    # from kafka.admin.new_topic import NewTopic
+
+    topic = "test-topic-%s" % str(uuid.uuid4())
+
+    # admin = KafkaAdminClient(bootstrap_servers=BROKER)
+    # new_topics = [NewTopic(topic, num_partitions=1, replication_factor=1)]
+    # topics = admin.create_topics(new_topics)
+
+    yield topic
+
+    # admin.delete_topics([topic])
 
 
 @pytest.fixture(scope="session")
@@ -96,7 +112,7 @@ def data_source():
 
 @transient_function_wrapper(kafka.producer.kafka, "KafkaProducer.send.__wrapped__")
 # Place transient wrapper underneath instrumentation
-def cache_kafka_headers(wrapped, instance, args, kwargs):
+def cache_kafka_producer_headers(wrapped, instance, args, kwargs):
     transaction = current_transaction()
 
     if transaction is None:
@@ -107,3 +123,29 @@ def cache_kafka_headers(wrapped, instance, args, kwargs):
     headers = dict(headers)
     transaction._test_request_headers = headers
     return ret
+
+
+@transient_function_wrapper(kafka.consumer.group, "KafkaConsumer.__next__")
+# Place transient wrapper underneath instrumentation
+def cache_kafka_consumer_headers(wrapped, instance, args, kwargs):
+    record = wrapped(*args, **kwargs)
+    transaction = current_transaction()
+
+    if transaction is None:
+        return record
+
+    headers = record.headers
+    headers = dict(headers)
+    transaction._test_request_headers = headers
+    return record
+
+
+@pytest.fixture(autouse=True)
+def assert_no_active_transaction():
+    # Run before test
+    assert not current_transaction(active_only=False), "Transaction exists before test run."
+
+    yield  # Run test
+
+    # Run after test
+    assert not current_transaction(active_only=False), "Transaction was not properly exited."

tests/messagebroker_kafkapython/test_kafka_consumer.py renamed to tests/messagebroker_kafkapython/test_consumer.py

@@ -12,21 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import kafka
 import kafka.errors as Errors
 import pytest
+from conftest import BROKER, cache_kafka_consumer_headers
 from testing_support.fixtures import (
     validate_attributes,
     validate_error_event_attributes_outside_transaction,
     validate_transaction_errors,
     validate_transaction_metrics,
 )
+from testing_support.validators.validate_distributed_trace_accepted import (
+    validate_distributed_trace_accepted,
+)
 
 from newrelic.api.background_task import background_task
 from newrelic.api.transaction import end_of_transaction
 from newrelic.packages import six
 
 
-def test_custom_metrics_are_recorded(get_consumer_records, topic):
+def test_custom_metrics(get_consumer_records, topic):
     @validate_transaction_metrics(
         "Named/%s" % topic,
         group="Message/Kafka/Topic",
@@ -42,11 +47,9 @@ def _test():
     _test()
 
 
-def test_custom_metrics_are_recorded_on_already_active_transaction(get_consumer_records, topic):
+def test_custom_metrics_on_existing_transaction(get_consumer_records, topic):
     transaction_name = (
-        "test_kafka_consumer:test_custom_metrics_are_recorded_on_already_active_transaction.<locals>._test"
-        if six.PY3
-        else "test_kafka_consumer:_test"
+        "test_consumer:test_custom_metrics_on_existing_transaction.<locals>._test" if six.PY3 else "test_consumer:_test"
     )
 
     @validate_transaction_metrics(
@@ -64,11 +67,9 @@ def _test():
     _test()
 
 
-def test_custom_metrics_are_not_recorded_on_inactive_transaction(get_consumer_records, topic):
+def test_custom_metrics_inactive_transaction(get_consumer_records, topic):
     transaction_name = (
-        "test_kafka_consumer:test_custom_metrics_are_not_recorded_on_inactive_transaction.<locals>._test"
-        if six.PY3
-        else "test_kafka_consumer:_test"
+        "test_consumer:test_custom_metrics_inactive_transaction.<locals>._test" if six.PY3 else "test_consumer:_test"
     )
 
     @validate_transaction_metrics(
@@ -87,15 +88,15 @@ def _test():
     _test()
 
 
-def test_agent_attributes_are_recorded(get_consumer_records):
+def test_agent_attributes(get_consumer_records):
     @validate_attributes("agent", ["kafka.consume.client_id", "kafka.consume.byteCount"])
     def _test():
         get_consumer_records()
 
     _test()
 
 
-def test_agent_records_error_if_raised(get_consumer_records, consumer_next_raises):
+def test_consumer_errors(get_consumer_records, consumer_next_raises):
     @validate_error_event_attributes_outside_transaction(
         exact_attrs={"intrinsic": {"error.class": "kafka.errors:KafkaError"}}
     )
@@ -106,7 +107,23 @@ def _test():
     _test()
 
 
-def test_agent_does_not_record_error_if_not_raised(get_consumer_records):
+def test_consumer_deserialization_errors(topic, consumer):
+    producer = kafka.KafkaProducer(
+        bootstrap_servers=BROKER, api_version=(2, 0, 2), value_serializer=lambda v: str(v).encode("utf-8")
+    )  # Producer that allows us to upload invalid JSON.
+
+    @validate_error_event_attributes_outside_transaction(exact_attrs={"intrinsic": {"error.class": "ValueError"}})
+    def _test():
+        with pytest.raises(ValueError):
+            producer.send(topic, value="%")  # Invalid JSON
+            producer.flush()
+            for _ in consumer:
+                pass
+
+    _test()
+
+
+def test_consumer_handled_errors_not_recorded(get_consumer_records):
     # It's important to check that we do not notice the StopIteration error.
     @validate_transaction_errors([])
     def _test():
@@ -115,6 +132,41 @@ def _test():
     _test()
 
 
+def test_distributed_tracing_headers(topic, producer, consumer):
+    # Send the messages inside a transaction, making sure to close it.
+    @background_task()
+    def _produce():
+        producer.send(topic, value={"foo": "bar"})
+        producer.flush()
+
+    consumer_iter = iter(consumer)
+
+    @validate_transaction_metrics(
+        "Named/%s" % topic,
+        group="Message/Kafka/Topic",
+        rollup_metrics=[
+            ("Supportability/DistributedTrace/AcceptPayload/Success", None),
+            ("Supportability/TraceContext/Accept/Success", 1),
+        ],
+        background_task=True,
+    )
+    def _consume():
+        @validate_distributed_trace_accepted(transport_type="Kafka")
+        @cache_kafka_consumer_headers
+        def _test():
+            # Start the transaction but don't exit it.
+            next(consumer_iter)
+
+        _test()
+
+        # Exit the transaction.
+        with pytest.raises(StopIteration):
+            next(consumer_iter)
+
+    _produce()
+    _consume()
+
+
 @pytest.fixture()
 def get_consumer_records(topic, producer, consumer):
     def _test():
@@ -132,4 +184,4 @@ def _poll(*args, **kwargs):
         raise Errors.KafkaError()
 
     consumer.poll = _poll
-    consumer
+    return consumer

tests/messagebroker_kafkapython/test_metrics.py

@@ -12,112 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from newrelic.packages import six
+
 
 def test_data_source_metrics(data_source, topic, producer, consumer):
+    _data_source_metrics = {
+        "MessageBroker/Kafka/Internal/kafka-metrics-count/count": "present",
+        "MessageBroker/Kafka/Internal/producer-metrics/request-rate": "present",
+        "MessageBroker/Kafka/Internal/producer-topic-metrics.%s/record-send-rate" % topic: "present",
+        "MessageBroker/Kafka/Internal/consumer-metrics/request-rate": "present",
+    }
+
     producer.send(topic, value=1)
     producer.flush()
-    next(iter(consumer))
+    for _ in consumer:
+        pass
 
     metrics = dict(data_source())
-    metric_names = list(metrics.keys())
     assert metrics
 
-
-# Example metrics
-
-# MessageBroker/Kafka/Internal/kafka-metrics-count/count
-# MessageBroker/Kafka/Internal/producer-metrics/connection-close-rate
-# MessageBroker/Kafka/Internal/producer-metrics/connection-creation-rate
-# MessageBroker/Kafka/Internal/producer-metrics/select-rate
-# MessageBroker/Kafka/Internal/producer-metrics/io-wait-time-ns-avg
-# MessageBroker/Kafka/Internal/producer-metrics/io-wait-ratio
-# MessageBroker/Kafka/Internal/producer-metrics/io-time-ns-avg
-# MessageBroker/Kafka/Internal/producer-metrics/io-ratio
-# MessageBroker/Kafka/Internal/producer-metrics/connection-count
-# MessageBroker/Kafka/Internal/producer-metrics/batch-size-avg
-# MessageBroker/Kafka/Internal/producer-metrics/batch-size-max
-# MessageBroker/Kafka/Internal/producer-metrics/compression-rate-avg
-# MessageBroker/Kafka/Internal/producer-metrics/record-queue-time-avg
-# MessageBroker/Kafka/Internal/producer-metrics/record-queue-time-max
-# MessageBroker/Kafka/Internal/producer-metrics/record-send-rate
-# MessageBroker/Kafka/Internal/producer-metrics/records-per-request-avg
-# MessageBroker/Kafka/Internal/producer-metrics/byte-rate
-# MessageBroker/Kafka/Internal/producer-metrics/record-size-max
-# MessageBroker/Kafka/Internal/producer-metrics/record-size-avg
-# MessageBroker/Kafka/Internal/producer-metrics/metadata-age
-# MessageBroker/Kafka/Internal/producer-metrics/network-io-rate
-# MessageBroker/Kafka/Internal/producer-metrics/outgoing-byte-rate
-# MessageBroker/Kafka/Internal/producer-metrics/request-rate
-# MessageBroker/Kafka/Internal/producer-metrics/request-size-avg
-# MessageBroker/Kafka/Internal/producer-metrics/request-size-max
-# MessageBroker/Kafka/Internal/producer-metrics/incoming-byte-rate
-# MessageBroker/Kafka/Internal/producer-metrics/response-rate
-# MessageBroker/Kafka/Internal/producer-metrics/request-latency-avg
-# MessageBroker/Kafka/Internal/producer-metrics/request-latency-max
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/outgoing-byte-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/request-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/request-size-avg
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/request-size-max
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/incoming-byte-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/response-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/request-latency-avg
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-bootstrap-0/request-latency-max
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/outgoing-byte-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/request-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/request-size-avg
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/request-size-max
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/incoming-byte-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/response-rate
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/request-latency-avg
-# MessageBroker/Kafka/Internal/producer-node-metrics.node-1001/request-latency-max
-# MessageBroker/Kafka/Internal/producer-topic-metrics.test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/record-send-rate
-# MessageBroker/Kafka/Internal/producer-topic-metrics.test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/byte-rate
-# MessageBroker/Kafka/Internal/producer-topic-metrics.test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/compression-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/connection-close-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/connection-creation-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/select-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/io-wait-time-ns-avg
-# MessageBroker/Kafka/Internal/consumer-metrics/io-wait-ratio
-# MessageBroker/Kafka/Internal/consumer-metrics/io-time-ns-avg
-# MessageBroker/Kafka/Internal/consumer-metrics/io-ratio
-# MessageBroker/Kafka/Internal/consumer-metrics/connection-count
-# MessageBroker/Kafka/Internal/consumer-metrics/network-io-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/outgoing-byte-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/request-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/request-size-avg
-# MessageBroker/Kafka/Internal/consumer-metrics/request-size-max
-# MessageBroker/Kafka/Internal/consumer-metrics/incoming-byte-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/response-rate
-# MessageBroker/Kafka/Internal/consumer-metrics/request-latency-avg
-# MessageBroker/Kafka/Internal/consumer-metrics/request-latency-max
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/outgoing-byte-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/request-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/request-size-avg
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/request-size-max
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/incoming-byte-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/response-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/request-latency-avg
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-bootstrap-0/request-latency-max
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/fetch-size-avg
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/fetch-size-max
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/bytes-consumed-rate
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/records-per-request-avg
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/records-consumed-rate
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/fetch-latency-avg
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/fetch-latency-max
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/fetch-rate
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/records-lag-max
-# MessageBroker/Kafka/Internal/consumer-coordinator-metrics/assigned-partitions
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/outgoing-byte-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/request-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/request-size-avg
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/request-size-max
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/incoming-byte-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/response-rate
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/request-latency-avg
-# MessageBroker/Kafka/Internal/consumer-node-metrics.node-1001/request-latency-max
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/fetch-size-avg
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/fetch-size-max
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/bytes-consumed-rate
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/records-per-request-avg
-# MessageBroker/Kafka/Internal/consumer-fetch-manager-metrics/test-topic-c962647a-f6cf-4a24-a90b-40ab2364dc55/records-consumed-rate
+    for metric_name, count in six.iteritems(_data_source_metrics):
+        if count == "present":
+            assert metric_name in metrics
+        else:
+            assert metrics[metric_name]["count"] == count, "%s:%d" % (metric_name, count)