fix: psycopg instrumentation fixes (#29)

Author: sohankshirsagar

drift/instrumentation/django/e2e-tests/src/test_requests.py

@@ -1,26 +1,6 @@
 """Execute test requests against the Django app."""
 
-import os
-import time
-
-import requests
-
-PORT = os.getenv("PORT", "8000")
-BASE_URL = f"http://localhost:{PORT}"
-
-
-def make_request(method: str, endpoint: str, **kwargs):
-    """Make HTTP request and log result."""
-    url = f"{BASE_URL}{endpoint}"
-    print(f"→ {method} {endpoint}")
-
-    # Set default timeout if not provided
-    kwargs.setdefault("timeout", 30)
-    response = requests.request(method, url, **kwargs)
-    print(f"  Status: {response.status_code}")
-    time.sleep(0.5)  # Small delay between requests
-    return response
-
+from drift.instrumentation.e2e_common.test_utils import make_request, print_request_summary
 
 if __name__ == "__main__":
     print("Starting Django test request sequence...\n")
@@ -42,4 +22,4 @@ def make_request(method: str, endpoint: str, **kwargs):
     )
     make_request("DELETE", "/api/post/1/delete")
 
-    print("\nAll requests completed successfully")
+    print_request_summary()

drift/instrumentation/e2e_common/__init__.py

@@ -1,5 +1,12 @@
 """Common utilities for Python SDK e2e tests."""
 
 from .base_runner import Colors, E2ETestRunnerBase
+from .test_utils import get_request_count, make_request, print_request_summary
 
-__all__ = ["Colors", "E2ETestRunnerBase"]
+__all__ = [
+    "Colors",
+    "E2ETestRunnerBase",
+    "get_request_count",
+    "make_request",
+    "print_request_summary",
+]

drift/instrumentation/e2e_common/base_runner.py

@@ -45,6 +45,7 @@ def __init__(self, app_port: int = 8000):
         self.app_port = app_port
         self.app_process: subprocess.Popen | None = None
         self.exit_code = 0
+        self.expected_request_count: int | None = None
 
         # Register signal handlers for cleanup
         signal.signal(signal.SIGTERM, self._signal_handler)
@@ -76,6 +77,17 @@ def run_command(self, cmd: list[str], env: dict | None = None, check: bool = Tru
 
         return result
 
+    def _parse_request_count(self, output: str):
+        """Parse the request count from test_requests.py output."""
+        for line in output.split("\n"):
+            if line.startswith("TOTAL_REQUESTS_SENT:"):
+                try:
+                    count = int(line.split(":")[1])
+                    self.expected_request_count = count
+                    self.log(f"Captured request count: {count}", Colors.GREEN)
+                except (ValueError, IndexError):
+                    self.log(f"Failed to parse request count from: {line}", Colors.YELLOW)
+
     def wait_for_service(self, check_cmd: list[str], timeout: int = 30, interval: int = 1) -> bool:
         """Wait for a service to become ready."""
         elapsed = 0
@@ -138,8 +150,8 @@ def record_traces(self) -> bool:
         self.app_process = subprocess.Popen(
             ["python", "src/app.py"],
             env={**os.environ, **env},
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
             text=True,
         )
 
@@ -164,7 +176,13 @@ def record_traces(self) -> bool:
         # Execute test requests
         self.log("Executing test requests...", Colors.GREEN)
         try:
-            self.run_command(["python", "src/test_requests.py"])
+            # Pass PYTHONPATH so test_requests.py can import from e2e_common
+            result = self.run_command(
+                ["python", "src/test_requests.py"],
+                env={"PYTHONPATH": "/sdk"},
+            )
+            # Parse request count from output
+            self._parse_request_count(result.stdout)
         except subprocess.CalledProcessError:
             self.log("Test requests failed", Colors.RED)
             self.exit_code = 1
@@ -257,13 +275,15 @@ def parse_test_results(self, output: str):
                     idx += 1
 
             all_passed = True
+            passed_count = 0
             for result in results:
                 test_id = result.get("test_id", "unknown")
                 passed = result.get("passed", False)
                 duration = result.get("duration", 0)
 
                 if passed:
                     self.log(f"✓ Test ID: {test_id} (Duration: {duration}ms)", Colors.GREEN)
+                    passed_count += 1
                 else:
                     self.log(f"✗ Test ID: {test_id} (Duration: {duration}ms)", Colors.RED)
                     all_passed = False
@@ -276,6 +296,20 @@ def parse_test_results(self, output: str):
                 self.log("Some tests failed!", Colors.RED)
                 self.exit_code = 1
 
+            # Validate request count matches passed tests
+            if self.expected_request_count is not None:
+                if passed_count < self.expected_request_count:
+                    self.log(
+                        f"✗ Request count mismatch: {passed_count} passed tests != {self.expected_request_count} requests sent",
+                        Colors.RED,
+                    )
+                    self.exit_code = 1
+                else:
+                    self.log(
+                        f"✓ Request count validation: {passed_count} passed tests >= {self.expected_request_count} requests sent",
+                        Colors.GREEN,
+                    )
+
         except Exception as e:
             self.log(f"Failed to parse test results: {e}", Colors.RED)
             self.log(f"Raw output:\n{output}", Colors.YELLOW)

drift/instrumentation/e2e_common/test_utils.py

@@ -0,0 +1,37 @@
+"""
+Shared test utilities for e2e tests.
+
+This module provides common functions used across all instrumentation e2e tests,
+including request counting for validation.
+"""
+
+import time
+
+import requests
+
+BASE_URL = "http://localhost:8000"
+_request_count = 0
+
+
+def make_request(method, endpoint, **kwargs):
+    """Make HTTP request, log result, and track count."""
+    global _request_count
+    _request_count += 1
+
+    url = f"{BASE_URL}{endpoint}"
+    print(f"→ {method} {endpoint}")
+    kwargs.setdefault("timeout", 30)
+    response = requests.request(method, url, **kwargs)
+    print(f"  Status: {response.status_code}")
+    time.sleep(0.5)
+    return response
+
+
+def get_request_count():
+    """Return the current request count."""
+    return _request_count
+
+
+def print_request_summary():
+    """Print the total request count in a parseable format."""
+    print(f"\nTOTAL_REQUESTS_SENT:{_request_count}")

drift/instrumentation/fastapi/e2e-tests/src/test_requests.py

@@ -1,27 +1,9 @@
 """Execute test requests against the FastAPI app."""
 
 import json
-import os
-import time
 from pathlib import Path
 
-import requests
-
-PORT = os.getenv("PORT", "8000")
-BASE_URL = f"http://localhost:{PORT}"
-
-
-def make_request(method: str, endpoint: str, **kwargs):
-    """Make HTTP request and log result."""
-    url = f"{BASE_URL}{endpoint}"
-    print(f"→ {method} {endpoint}")
-
-    # Set default timeout if not provided
-    kwargs.setdefault("timeout", 30)
-    response = requests.request(method, url, **kwargs)
-    print(f"  Status: {response.status_code}")
-    time.sleep(0.5)  # Small delay between requests
-    return response
+from drift.instrumentation.e2e_common.test_utils import make_request, print_request_summary
 
 
 def verify_stack_traces():
@@ -180,3 +162,5 @@ def verify_context_propagation():
     print("\n" + "=" * 50)
     print("All requests completed successfully")
     print("=" * 50)
+
+    print_request_summary()

drift/instrumentation/flask/e2e-tests/src/test_requests.py

@@ -1,24 +1,6 @@
 """Execute test requests against the Flask app."""
 
-import time
-
-import requests
-
-BASE_URL = "http://localhost:8000"
-
-
-def make_request(method, endpoint, **kwargs):
-    """Make HTTP request and log result."""
-    url = f"{BASE_URL}{endpoint}"
-    print(f"→ {method} {endpoint}")
-
-    # Set default timeout if not provided
-    kwargs.setdefault("timeout", 30)
-    response = requests.request(method, url, **kwargs)
-    print(f"  Status: {response.status_code}")
-    time.sleep(0.5)  # Small delay between requests
-    return response
-
+from drift.instrumentation.e2e_common.test_utils import make_request, print_request_summary
 
 if __name__ == "__main__":
     print("Starting test request sequence...\n")
@@ -32,4 +14,4 @@ def make_request(method, endpoint, **kwargs):
     make_request("POST", "/api/post", json={"title": "Test Post", "body": "This is a test post", "userId": 1})
     make_request("DELETE", "/api/post/1")
 
-    print("\nAll requests completed successfully")
+    print_request_summary()

drift/instrumentation/httpx/e2e-tests/src/test_requests.py

@@ -1,24 +1,6 @@
 """Execute test requests against the Flask app to exercise the HTTPX instrumentation."""
 
-import time
-
-import requests
-
-BASE_URL = "http://localhost:8000"
-
-
-def make_request(method, endpoint, **kwargs):
-    """Make HTTP request and log result."""
-    url = f"{BASE_URL}{endpoint}"
-    print(f"-> {method} {endpoint}")
-
-    # Set default timeout if not provided
-    kwargs.setdefault("timeout", 30)
-    response = requests.request(method, url, **kwargs)
-    print(f"   Status: {response.status_code}")
-    time.sleep(0.5)  # Small delay between requests
-    return response
-
+from drift.instrumentation.e2e_common.test_utils import make_request, print_request_summary
 
 if __name__ == "__main__":
     print("Starting test request sequence for HTTPX instrumentation...\n")
@@ -124,4 +106,4 @@ def make_request(method, endpoint, **kwargs):
 
     make_request("POST", "/test/file-like-body")
 
-    print("\nAll requests completed successfully")
+    print_request_summary()

drift/instrumentation/httpx/instrumentation.py

@@ -823,6 +823,7 @@ def _try_get_mock_from_request_sync(
                 input_value=input_value,
                 kind=SpanKind.CLIENT,
                 input_schema_merges=input_schema_merges,
+                is_pre_app_start=not sdk.app_ready,
             )
 
             if not mock_response_output or not mock_response_output.found:

drift/instrumentation/psycopg/e2e-tests/requirements.txt

@@ -1,4 +1,5 @@
 -e /sdk  # Mount point for drift SDK
 Flask>=3.1.2
 psycopg[binary]>=3.2.1
+psycopg-pool>=3.2.0
 requests>=2.32.5