fix: redis span alignment (#474)

* fix: redis span alignment

* chore: extract exporter wait code

src/lumigo_opentelemetry/instrumentations/redis/__init__.py

@@ -1,7 +1,7 @@
-import json
-from typing import List, Any, Dict
+from typing import Any, Dict, List
 
 from opentelemetry.trace.span import Span
+
 from lumigo_opentelemetry.instrumentations import AbstractInstrumentor
 from lumigo_opentelemetry.instrumentations.instrumentation_utils import (
     add_body_attribute,
@@ -22,8 +22,9 @@ def install_instrumentation(self) -> None:
         def request_hook(
             span: Span, instance: Connection, args: List[Any], kwargs: Dict[Any, Any]
         ) -> None:
-            add_body_attribute(span, json.dumps(args), "redis.request.args")
-            add_body_attribute(span, json.dumps(kwargs), "redis.request.kwargs")
+            # a db.statement attribute is automatically added by the RedisInstrumentor
+            # when this hook is called, so we don't need to add anything manually
+            pass
 
         def response_hook(span: Span, instance: Connection, response: Any) -> None:
             add_body_attribute(span, response, "redis.response.body")

src/test/components/tests/test_attr_max_size.py

@@ -1,5 +1,5 @@
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -14,8 +14,7 @@ def test_large_span_attribute_size_max_size_env_var_was_set(self):
 
         assert body is not None
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(4, len(spans_container.spans))

src/test/components/tests/test_execution_tags.py

@@ -1,5 +1,5 @@
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -14,8 +14,7 @@ def test_execution_tag(self):
 
         assert body is not None
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(4, len(spans_container.spans))

src/test/integration/boto3/tests/test_boto3.py

@@ -1,6 +1,6 @@
 import json
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -35,8 +35,7 @@ def test_boto3_instrumentation(self):
 
         self.assertEqual(body, {"status": "ok"})
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(4, len(spans_container.spans))

src/test/integration/fastapi/tests/test_fastapi.py

@@ -1,5 +1,5 @@
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -14,8 +14,7 @@ def test_200_OK(self):
 
         self.assertEqual(body, {"message": "Hello FastAPI!"})
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(3, len(spans_container.spans))
@@ -50,8 +49,7 @@ def test_requests_instrumentation(self):
 
         self.assertIn("https://api.chucknorris.io/jokes/", body["url"])
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(4, len(spans_container.spans))
@@ -86,8 +84,7 @@ def test_large_span_attribute_size_default_max_size(self):
 
         assert body is not None
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(4, len(spans_container.spans))

src/test/integration/flask/tests/test_flask.py

@@ -1,5 +1,5 @@
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -14,8 +14,7 @@ def test_200_OK(self):
 
         self.assertEqual(body, {"message": "Hello Flask!"})
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(1, len(spans_container.spans))
@@ -37,8 +36,7 @@ def test_requests_instrumentation(self):
 
         self.assertIn("https://api.chucknorris.io/jokes/", body["url"])
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
         self.assertEqual(2, len(spans_container.spans))

src/test/integration/grpcio/tests/test_grpcio.py

@@ -1,7 +1,7 @@
-import time
-import unittest
-import sys
 import os
+import sys
+import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 
 sys.path.append(os.path.join(os.path.dirname(__file__), "..", "app"))
 
@@ -20,8 +20,7 @@ def tearDownClass(cls) -> None:
             stub.SayHelloUnaryUnary(helloworld_pb2.HelloRequest(name="exit"))
 
     def check_spans(self, method: str, request_payload: str, response_payload: str):
-        # TODO Do something deterministic
-        time.sleep(5)  # allow the exporter to catch up
+        wait_for_exporter(5)
 
         server_file = os.getenv("SERVER_SPANDUMP")
         server_spans = SpansContainer.get_spans_from_file(server_file)

src/test/integration/kafka_python/tests/test_kafka_python.py

@@ -1,6 +1,6 @@
 import json
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -43,8 +43,7 @@ def test_kafka_python_instrumentation(self):
 
             self.assertEqual(body, {"status": "ok"})
 
-            # TODO Do something deterministic
-            time.sleep(3)  # Sleep to allow the exporter to catch up
+            wait_for_exporter()
 
             spans_container = SpansContainer.get_spans_from_file()
 

src/test/integration/pika/app/__init__.py

@@ -46,6 +46,7 @@ async def invoke_pika_consumer(request: Request):
     except Exception as e:
         print(f"Exception while consuming messages: {e}")
         channel.stop_consuming()
+        raise e
     return {"status": "ok"}
 
 

src/test/integration/pika/tests/test_pika.py

@@ -1,6 +1,6 @@
 import json
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -48,8 +48,7 @@ def test_pika_instrumentation(self):
 
             self.assertEqual(body, {"status": "ok"})
 
-            # TODO Do something deterministic
-            time.sleep(3)  # Sleep to allow the exporter to catch up
+            wait_for_exporter()
 
             spans_container = SpansContainer.get_spans_from_file()
 

src/test/integration/pymongo/tests/test_pymongo.py

@@ -1,5 +1,5 @@
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -15,8 +15,7 @@ def test_mongo_instrumentation(self):
 
         self.assertEqual(body, {"status": "ok"})
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
 

src/test/integration/pymysql/tests/test_pymysql.py

@@ -1,5 +1,5 @@
-import time
 import unittest
+from test.test_utils.span_exporter import wait_for_exporter
 from test.test_utils.spans_parser import SpansContainer
 
 import requests
@@ -15,8 +15,7 @@ def test_pymysql_instrumentation(self):
 
         self.assertEqual(body, {"status": "ok"})
 
-        # TODO Do something deterministic
-        time.sleep(3)  # Sleep to allow the exporter to catch up
+        wait_for_exporter()
 
         spans_container = SpansContainer.get_spans_from_file()
 

src/test/integration/redis/app/redis_hash.py

@@ -0,0 +1,7 @@
+import os
+
+import redis
+
+client = redis.StrictRedis(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT"))
+client.hmset("my-key", {"key1": "value1", "key2": "value2"})
+client.hgetall("my-key")

src/test/integration/redis/app/redis_transaction.py

@@ -0,0 +1,10 @@
+import os
+
+import redis
+
+client = redis.StrictRedis(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT"))
+pipe = client.pipeline()
+pipe.set("my-key", "pre-key-value").get("my-key")
+pipe.set("my-key", "key-value").get("my-key")
+pipe.get("unknown-key")
+pipe.execute()