added support for self.callback_on_request and self.callback_on_response

owasp-sbot · Sep 12, 2024 · 2024c2a · 2024c2a
1 parent c2355e9
commit 2024c2a
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 16 deletions.
diff --git a/osbot_fast_api/api/Fast_API__Http_Events.py b/osbot_fast_api/api/Fast_API__Http_Events.py
@@ -1,3 +1,4 @@
+import types
 from collections                                    import deque
 from osbot_fast_api.api.Fast_API__Request_Data      import Fast_API__Request_Data
 from osbot_utils.base_classes.Type_Safe             import Type_Safe
@@ -12,6 +13,8 @@
 class Fast_API__Http_Events(Type_Safe):
     #log_requests          : bool = False                           # todo: change this to save on S3 and disk
     background_tasks      : list
+    callback_on_request   : types.MethodType
+    callback_on_response  : types.MethodType
     trace_calls           : bool = False
     trace_call_config     : Trace_Call__Config
     requests_data         : dict
@@ -27,16 +30,17 @@ def __init__(self,**kwargs):
     def on_http_request(self, request: Request):
         with self.request_data(request) as _:
             _.on_request(request)
-            #_.add_log_message("on_http_request")
-        self.request_trace_start(request)
+            self.request_trace_start(request)
+            if self.callback_on_request:
+                self.callback_on_request(_)
 
     def on_http_response(self, request: Request, response:Response):
         with self.request_data(request) as _:
             _.on_response(response)
-            #_.add_log_message("on_http_response")
-
-        # if StreamingResponse not in base_types(response):                          # handle the special case when the response is a StreamingResponse
-        self.request_trace_stop(request)                                             # todo: change this to be on text/event-stream"; charset=utf-8 (which is the one that happens with the LLMs responses)
+            # if StreamingResponse not in base_types(response):                          # handle the special case when the response is a StreamingResponse
+            self.request_trace_stop(request)                                             # todo: change this to be on text/event-stream"; charset=utf-8 (which is the one that happens with the LLMs responses)
+            if self.callback_on_response:
+                self.callback_on_response(_)
 
 
 

diff --git a/tests/unit/api/test_Fast_API__Http_Events.py b/tests/unit/api/test_Fast_API__Http_Events.py
@@ -39,6 +39,8 @@ def test__init__(self):
         with self.http_events as _:
 
             expected_locals = {'background_tasks'     : []                                   ,
+                               'callback_on_request'  : None                                 ,
+                               'callback_on_response' : None                                 ,
                                'fast_api_name'        : ''                                   ,
                                'requests_data'        : { self.request_id: self.request_data},
                                'requests_order'       : deque([self.request_id])             ,

diff --git a/tests/unit/utils/test_Uvicorn_Server.py b/tests/unit/utils/test_Uvicorn_Server.py
@@ -44,13 +44,12 @@ def test_read_stdout(self, mock_print):
         mock_print.assert_any_call('stdout:', 'test output\n', end='')
 
     def test_start_stop(self):
-        with Duration(prefix='start'):
-            assert is_port_open(UVICORN_SERVER_NAME, self.uvicorn_server.port) is False
-            assert self.uvicorn_server.start() is True
-
-        with Duration(prefix='requesrt'):
-            assert is_port_open(UVICORN_SERVER_NAME, self.uvicorn_server.port) is True
-            assert '<title>FastAPI - Swagger UI</title>' in self.uvicorn_server.http_GET('docs')
-        with Duration(prefix='stop'):
-            assert self.uvicorn_server.stop() is True
-            assert is_port_open(UVICORN_SERVER_NAME, self.uvicorn_server.port) is False
+
+        assert is_port_open(UVICORN_SERVER_NAME, self.uvicorn_server.port) is False
+        assert self.uvicorn_server.start() is True
+
+        assert is_port_open(UVICORN_SERVER_NAME, self.uvicorn_server.port) is True
+        assert '<title>FastAPI - Swagger UI</title>' in self.uvicorn_server.http_GET('docs')
+
+        assert self.uvicorn_server.stop() is True
+        assert is_port_open(UVICORN_SERVER_NAME, self.uvicorn_server.port) is False