dispatch response hooks in replay mode

Author: sohankshirsagar

drift/instrumentation/requests/e2e-tests/src/app.py

@@ -304,17 +304,9 @@ def test_streaming_iter_lines():
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-
-
-# ============ BUG DETECTION TESTS ============
-# These tests expose bugs in the requests instrumentation during REPLAY mode.
-# Keep only tests that expose confirmed bugs.
-
-# BUG 4: Response hooks are not called in REPLAY mode
-# dispatch_hook() is not called when returning mock response
 @app.route("/test/response-hooks", methods=["GET"])
 def test_response_hooks():
-    """Test response hooks that modify response - EXPOSES BUG."""
+    """Test response hooks that modify response."""
     try:
         hook_called = {"value": False}
         hook_status = {"value": 0}

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

@@ -74,11 +74,6 @@ def make_request(method, endpoint, **kwargs):
 
     make_request("GET", "/test/streaming-iter-lines")
 
-    # ============ BUG DETECTION TESTS ============
-    # These tests expose confirmed bugs in the requests instrumentation
-    print("\n=== Bug Detection Tests ===\n")
-
-    # BUG 4: Response hooks are not called in REPLAY mode
     make_request("GET", "/test/response-hooks")
 
     print("\nAll requests completed successfully")

drift/instrumentation/requests/instrumentation.py

@@ -132,7 +132,7 @@ def original_call():
             if sdk.mode == TuskDriftMode.REPLAY:
                 return handle_replay_mode(
                     replay_mode_handler=lambda: instrumentation_self._handle_replay_send(
-                        sdk, request
+                        sdk, request, **kwargs
                     ),
                     no_op_request_handler=lambda: instrumentation_self._get_default_response(url),
                     is_server_request=False,
@@ -264,17 +264,16 @@ def _handle_replay_send(
         self,
         sdk: TuskDrift,
         prepared_request: Any,
+        **kwargs,
     ) -> Any:
         """Handle send() in REPLAY mode.
 
         Similar to _handle_replay but works with PreparedRequest objects.
 
         Args:
             sdk: TuskDrift instance
-            session_self: Session instance
             prepared_request: PreparedRequest object
-            original_send: Original Session.send method
-            **kwargs: Additional send() kwargs
+            **kwargs: Additional send() kwargs (timeout, verify, cert, proxies, etc.)
         """
         method = prepared_request.method
         url = prepared_request.url
@@ -317,6 +316,10 @@ def _handle_replay_send(
                 )
 
                 if mock_response is not None:
+                    # Dispatch response hooks (matches Session.send() behavior)
+                    # This ensures hooks registered via hooks={"response": callback} are called
+                    from requests.hooks import dispatch_hook
+                    mock_response = dispatch_hook("response", prepared_request.hooks, mock_response, **kwargs)
                     return mock_response
 
                 # No mock found - raise error in REPLAY mode