add motor tests and payload to mongo (#512)

Author: saartochner-lumigo

noxfile.py

@@ -873,6 +873,49 @@ def integration_tests_pymongo(
                 kill_process_and_clean_outputs(temp_file, "uvicorn", session)
 
 
+@nox.session()
+@nox.parametrize(
+    "python,motor_version",
+    [
+        (python, motor_version)
+        for python in python_versions()
+        for motor_version in dependency_versions_to_be_tested(
+            python=python,
+            directory="motor",
+            dependency_name="motor",
+        )
+    ],
+)
+def integration_tests_motor(
+    session,
+    python,
+    motor_version,
+):
+    with TestedVersions.save_tests_result("motor", python, "motor", motor_version):
+        install_package("motor", motor_version, session)
+
+        session.install(".")
+
+        temp_file = create_it_tempfile("motor")
+        with session.chdir("src/test/integration/motor"):
+            session.install("-r", OTHER_REQUIREMENTS)
+            try:
+                session.run(
+                    "pytest",
+                    "--tb",
+                    "native",
+                    "--log-cli-level=INFO",
+                    "--color=yes",
+                    "-v",
+                    "./tests/test_motor.py",
+                    env={
+                        "LUMIGO_DEBUG_SPANDUMP": temp_file,
+                    },
+                )
+            finally:
+                kill_process_and_clean_outputs(temp_file, "test_motor", session)
+
+
 @nox.session()
 @nox.parametrize(
     "python,pymysql_version",

src/lumigo_opentelemetry/instrumentations/pymongo/__init__.py

@@ -1,4 +1,11 @@
+from typing import Union
+
+from opentelemetry.trace.span import Span
+
 from lumigo_opentelemetry.instrumentations import AbstractInstrumentor
+from lumigo_opentelemetry import logger
+from lumigo_opentelemetry.libs.general_utils import lumigo_safe_execute
+from lumigo_opentelemetry.libs.json_utils import dump_with_context
 
 
 class PymongoInstrumentor(AbstractInstrumentor):
@@ -10,8 +17,43 @@ def check_if_applicable(self) -> None:
 
     def install_instrumentation(self) -> None:
         from opentelemetry.instrumentation.pymongo import PymongoInstrumentor
+        from pymongo.monitoring import (
+            CommandStartedEvent,
+            CommandSucceededEvent,
+            CommandFailedEvent,
+        )
+
+        def request_hook(span: Span, event: CommandStartedEvent) -> None:
+            with lumigo_safe_execute("pymongo_request_hook"):
+                span.set_attribute("db.statement", event.command_name)
+                if isinstance(event, CommandStartedEvent):
+                    span.set_attribute(
+                        "db.request.body",
+                        dump_with_context("requestBody", event.command),
+                    )
+                else:
+                    logger.warning(f"Got unexpected event type {type(event)}")
+
+        def response_hook(
+            span: Span, event: Union[CommandSucceededEvent, CommandFailedEvent]
+        ) -> None:
+            with lumigo_safe_execute("pymongo_response_hook"):
+                if isinstance(event, CommandSucceededEvent):
+                    span.set_attribute(
+                        "db.response.body",
+                        dump_with_context("responseBody", event.reply),
+                    )
+                elif isinstance(event, CommandFailedEvent):
+                    span.set_attribute(
+                        "db.response.body",
+                        dump_with_context("responseBody", event.failure),
+                    )
+                else:
+                    logger.warning(f"Got unexpected event type {type(event)}")
 
-        PymongoInstrumentor().instrument()
+        PymongoInstrumentor().instrument(
+            request_hook=request_hook, response_hook=response_hook
+        )
 
 
 instrumentor: AbstractInstrumentor = PymongoInstrumentor()

src/test/integration/motor/__init__.py

@@ -0,0 +1,23 @@
+from testcontainers.mongodb import MongoDbContainer
+
+import motor.motor_asyncio
+
+
+async def do_insert():
+    document = {"my-k": "my-value"}
+    result = await db.test_collection.insert_one(document)
+    print(f"result {result.inserted_id}")
+
+
+async def do_find_one():
+    document = await db.test_collection.find_one()
+    print("found", document)
+
+
+with MongoDbContainer("mongo:4") as mongo:
+    client = motor.motor_asyncio.AsyncIOMotorClient(mongo.get_connection_url())
+    db = client.test
+
+    loop = client.get_io_loop()
+    loop.run_until_complete(do_insert())
+    loop.run_until_complete(do_find_one())

src/test/integration/motor/app/motor_set_and_get.py

@@ -0,0 +1,3 @@
+pytest==7.1.1
+psutil==5.9.1
+testcontainers==3.7.0

src/test/integration/motor/requirements_others.txt

@@ -0,0 +1 @@
+3.2.0

src/test/integration/motor/tested_versions/3.10/motor

@@ -0,0 +1 @@
+3.2.0

src/test/integration/motor/tested_versions/3.11/motor

@@ -0,0 +1 @@
+3.2.0

src/test/integration/motor/tested_versions/3.7/motor

@@ -0,0 +1 @@
+3.2.0

src/test/integration/motor/tested_versions/3.8/motor

@@ -0,0 +1 @@
+3.2.0

src/test/integration/motor/tested_versions/3.9/motor

@@ -0,0 +1,51 @@
+import os
+import subprocess
+import sys
+import unittest
+from test.test_utils.span_exporter import wait_for_exporter
+from test.test_utils.spans_parser import SpansContainer
+
+
+class TestMotorSpans(unittest.TestCase):
+    """
+    This package uses pymongo under the hood, so there's no specific instrumentation for it,
+        we only want to measure our compatibility with it.
+    """
+
+    def test_motor_set_and_get(self):
+        sample_path = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)),
+            "../app/motor_set_and_get.py",
+        )
+        subprocess.check_output(
+            [sys.executable, sample_path],
+            env={
+                **os.environ,
+                "AUTOWRAPT_BOOTSTRAP": "lumigo_opentelemetry",
+                "OTEL_SERVICE_NAME": "motor",
+            },
+        )
+
+        wait_for_exporter()
+
+        spans = [
+            s
+            for s in SpansContainer.parse_spans_from_file().spans
+            if s["attributes"].get("db.system") == "mongodb"
+        ]
+
+        self.assertEqual(len(spans), 3)
+        admin_span, set_span, get_span = spans
+        self.assertEqual(set_span["attributes"]["db.statement"], "insert")
+
+        self.assertTrue(
+            '"my-k": "my-value"' in set_span["attributes"]["db.request.body"]
+        )
+        self.assertTrue(
+            '{"n": 1, "ok": 1.0}' in set_span["attributes"]["db.response.body"]
+        )
+
+        self.assertEqual(get_span["attributes"]["db.statement"], "find")
+        self.assertTrue(
+            '"my-k": "my-value"' in get_span["attributes"]["db.response.body"]
+        )

src/test/integration/motor/tests/__init__.py

@@ -5,7 +5,7 @@
 import requests
 
 
-class TestFastApiSpans(unittest.TestCase):
+class TestMongoSpans(unittest.TestCase):
     def test_mongo_instrumentation(self):
         response = requests.get("http://localhost:8002/invoke-mongo")
 
@@ -17,23 +17,30 @@ def test_mongo_instrumentation(self):
 
         wait_for_exporter()
 
-        spans_container = SpansContainer.get_spans_from_file()
+        spans = [
+            s
+            for s in SpansContainer.parse_spans_from_file().spans
+            if s["attributes"].get("db.system") == "mongodb"
+        ]
+
+        set_span, get_span = spans[
+            -2:
+        ]  # there are more spans to connect to the mongo container
 
         # assert mongo children spans
-        assert spans_container.find_child_span(
-            lambda span: (
-                span["kind"] == "SpanKind.CLIENT"
-                and span["attributes"]["db.system"] == "mongodb"
-                and span["attributes"]["db.name"] == "test"
-                and span["attributes"]["db.statement"] == "insert items"
-            )
+        assert (
+            set_span["kind"] == "SpanKind.CLIENT"
+            and set_span["attributes"]["db.system"] == "mongodb"
+            and set_span["attributes"]["db.name"] == "test"
+            and set_span["attributes"]["db.statement"] == "insert"
+            and '"street": "2 Avenue"' in set_span["attributes"]["db.request.body"]
+            and '{"n": 1, "ok": 1.0}' in set_span["attributes"]["db.response.body"]
         )
 
-        assert spans_container.find_child_span(
-            lambda span: (
-                span["kind"] == "SpanKind.CLIENT"
-                and span["attributes"]["db.system"] == "mongodb"
-                and span["attributes"]["db.name"] == "test"
-                and span["attributes"]["db.statement"] == "find items"
-            )
+        assert (
+            get_span["kind"] == "SpanKind.CLIENT"
+            and get_span["attributes"]["db.system"] == "mongodb"
+            and get_span["attributes"]["db.name"] == "test"
+            and get_span["attributes"]["db.statement"] == "find"
+            and '"street": "2 Avenue"' in get_span["attributes"]["db.response.body"]
         )