Adds a test to make sure masking is working

open-telemetry · Jun 25, 2020 · a06bfea · a06bfea
1 parent 9f5e7dc
commit a06bfea
Showing 1 changed file with 45 additions and 3 deletions.
diff --git a/ext/opentelemetry-ext-docker-tests/tests/asyncpg/test_asyncpg_functional.py b/ext/opentelemetry-ext-docker-tests/tests/asyncpg/test_asyncpg_functional.py
@@ -1,4 +1,6 @@
 import asyncio
+import hashlib
+import json
 import os
 
 import asyncpg
@@ -195,15 +197,14 @@ def test_instrumented_method_doesnt_capture_parameters(self, *_, **__):
 
 
 
-
-class TestFunctionalAsyncPG_CaptureParameters(TestBase):
+class TestFunctionalAsyncPG_CaptureParametersUnmasked(TestBase):
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
         cls._connection = None
         cls._cursor = None
         cls._tracer = cls.tracer_provider.get_tracer(__name__)
-        AsyncPGInstrumentor(capture_parameters=True).instrument(tracer_provider=cls.tracer_provider)
+        AsyncPGInstrumentor(capture_parameters=True, mask_parameters=False).instrument(tracer_provider=cls.tracer_provider)
         cls._connection = _await(
             asyncpg.connect(
                 database=POSTGRES_DB_NAME,
@@ -284,4 +285,45 @@ def test_instrumented_execute_interface_error_method(self, *_, **__):
                 "db.statement.parameters": "(1, 2, 3)",
                 "db.statement": "SELECT 42;",
             },
+        )
+
+class TestFunctionalAsyncPG_CaptureParametersMasked(TestBase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls._connection = None
+        cls._cursor = None
+        cls._tracer = cls.tracer_provider.get_tracer(__name__)
+        AsyncPGInstrumentor(capture_parameters=True).instrument(tracer_provider=cls.tracer_provider)
+        cls._connection = _await(
+            asyncpg.connect(
+                database=POSTGRES_DB_NAME,
+                user=POSTGRES_USER,
+                password=POSTGRES_PASSWORD,
+                host=POSTGRES_HOST,
+                port=POSTGRES_PORT,
+            )
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        AsyncPGInstrumentor().uninstrument()
+
+    @pytest.mark.asyncpg
+    def test_instrumented_execute_method_with_arguments(self, *_, **__):
+        _await(self._connection.execute("SELECT $1;", "1"))
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        self.assertEqual(
+            StatusCanonicalCode.OK, spans[0].status.canonical_code
+        )
+        self.assertEqual(
+            spans[0].attributes,
+            {
+                "db.type": "sql",
+                "db.user": POSTGRES_USER,
+                "db.statement.parameters": "sha256:" + hashlib.sha256(json.dumps(('1',)).encode('utf-8')).hexdigest(),
+                "db.instance": POSTGRES_DB_NAME,
+                "db.statement": "SELECT $1;",
+            },
         )