instrumentation/kafka: fix handling consumer iteration if transaction…

… not sampled (#2075)

Handle the case where if the transaction is not sampled capture_span
will return None instead of span.
While at it fix handling of checking for KAFKA_HOST in tests.

Fix #2073

elasticapm/instrumentation/packages/kafka.py

@@ -143,7 +143,8 @@ def call(self, module, method, wrapped, instance, args, kwargs):
  try:
  result = wrapped(*args, **kwargs)
  except StopIteration:
- span.cancel()
+ if span:
+ span.cancel()
  raise
  if span and not isinstance(span, DroppedSpan):
  topic = result[0]

tests/instrumentation/kafka_tests.py

@@ -45,11 +45,10 @@
 
 pytestmark = [pytest.mark.kafka]
 
-if "KAFKA_HOST" not in os.environ:
+KAFKA_HOST = os.environ.get("KAFKA_HOST")
+if not KAFKA_HOST:
  pytestmark.append(pytest.mark.skip("Skipping kafka tests, no KAFKA_HOST environment variable set"))
 
-KAFKA_HOST = os.environ["KAFKA_HOST"]
-
 
 @pytest.fixture(scope="function")
 def topics():
@@ -233,3 +232,22 @@ def test_kafka_poll_unsampled_transaction(instrument, elasticapm_client, consume
  elasticapm_client.end_transaction("foo")
  spans = elasticapm_client.events[SPAN]
  assert len(spans) == 0
+
+
+def test_kafka_consumer_unsampled_transaction_handles_stop_iteration(
+ instrument, elasticapm_client, producer, consumer, topics
+):
+ def delayed_send():
+ time.sleep(0.2)
+ producer.send("test", key=b"foo", value=b"bar")
+
+ thread = threading.Thread(target=delayed_send)
+ thread.start()
+ transaction = elasticapm_client.begin_transaction("foo")
+ transaction.is_sampled = False
+ for item in consumer:
+ pass
+ thread.join()
+ elasticapm_client.end_transaction("foo")
+ spans = elasticapm_client.events[SPAN]
+ assert len(spans) == 0