Implement jwt fixture, centralized reusable mock_req into conftest.py

Signed-off-by: Teo <[email protected]>

reauthorize

Signed-off-by: Teo <[email protected]>

tests/unit/conftest.py

@@ -1,5 +1,7 @@
 import contextlib
-from typing import Iterator
+import uuid
+from collections import defaultdict
+from typing import Dict, Iterator, List
 
 import pytest
 import requests_mock
@@ -9,8 +11,20 @@
 from agentops.singleton import clear_singletons
 from tests.fixtures.event import llm_event_spy
 
-# Common JWT tokens used across tests
-JWTS = ["some_jwt", "some_jwt2", "some_jwt3"]
+
+@pytest.fixture
+def jwt():
+    """Fixture that provides unique JWTs per session within a test"""
+    session_jwts = defaultdict(lambda: str(uuid.uuid4()))
+    session_count = 0
+
+    def get_jwt():
+        nonlocal session_count
+        jwt = session_jwts[session_count]
+        session_count += 1
+        return jwt
+
+    return get_jwt
 
 
 @pytest.fixture(autouse=True)
@@ -23,38 +37,6 @@ def setup_teardown():
     agentops.end_all_sessions()  # teardown part
 
 
-# @contextlib.contextmanager
-# @pytest.fixture(autouse=True)
-# def mock_req() -> Iterator[requests_mock.Mocker]:
-#     """
-#     Centralized request mocking for all tests.
-#     Mocks common API endpoints with standard responses.
-#     """
-#     with requests_mock.Mocker() as m:
-#         url = "https://api.agentops.ai"
-#
-#         # Mock API endpoints
-#         m.post(url + "/v2/create_events", json={"status": "ok"})
-#         m.post(url + "/v2/developer_errors", json={"status": "ok"})
-#         m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-#         m.post("https://pypi.org/pypi/agentops/json", status_code=404)
-#
-#         # Use iterator for JWT tokens in session creation
-#         jwt_tokens = iter(JWTS)
-#
-#         def create_session_response(request, context):
-#             context.status_code = 200
-#             try:
-#                 return {"status": "success", "jwt": next(jwt_tokens)}
-#             except StopIteration:
-#                 return {"status": "success", "jwt": "some_jwt"}  # Fallback JWT
-#
-#         m.post(url + "/v2/create_session", json=create_session_response)
-#         m.post(url + "/v2/reauthorize_jwt", json={"status": "success", "jwt": "some_jwt"})
-#
-#         yield m
-
-
 @pytest.fixture(scope="session")
 def api_key() -> str:
     """Standard API key for testing"""
@@ -65,3 +47,38 @@ def api_key() -> str:
 def base_url() -> str:
     """Base API URL"""
     return "https://api.agentops.ai"
+
+
+@pytest.fixture(autouse=True)
+def mock_req(base_url, jwt):
+    """
+    Mocks AgentOps backend API requests.
+    """
+    with requests_mock.Mocker() as m:
+        # Map session IDs to their JWTs
+        m.session_jwts = {}
+
+        m.post(base_url + "/v2/create_events", json={"status": "ok"})
+
+        def create_session_response(request, context):
+            context.status_code = 200
+            # Extract session_id from the request
+            session_id = request.json()["session"]["session_id"]
+            # Use the jwt fixture to get consistent JWTs
+            m.session_jwts[session_id] = jwt()
+            return {"status": "success", "jwt": m.session_jwts[session_id]}
+
+        def reauthorize_jwt_response(request, context):
+            context.status_code = 200
+            # Extract session_id from the request
+            session_id = request.json()["session_id"]
+            # Return the same JWT for this session
+            return {"status": "success", "jwt": m.session_jwts[session_id]}
+
+        m.post(base_url + "/v2/create_session", json=create_session_response)
+        m.post(base_url + "/v2/update_session", json={"status": "success", "token_cost": 5})
+        m.post(base_url + "/v2/developer_errors", json={"status": "ok"})
+        m.post(base_url + "/v2/reauthorize_jwt", json=reauthorize_jwt_response)
+        m.post(base_url + "/v2/create_agent", json={"status": "success"})
+
+        yield m

tests/unit/test_canary.py

@@ -7,18 +7,6 @@
 from agentops import ActionEvent
 
 
-@pytest.fixture(autouse=True, scope="function")
-def mock_req():
-    with requests_mock.Mocker() as m:
-        url = "https://api.agentops.ai"
-        m.post(url + "/v2/create_events", json={"status": "ok"})
-        m.post(url + "/v2/create_session", json={"status": "success", "jwt": "some_jwt"})
-        m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", json={"status": "ok"})
-        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
-        yield m
-
-
 class TestCanary:
     def setup_method(self):
         self.url = "https://api.agentops.ai"

tests/unit/test_events.py

@@ -7,20 +7,6 @@
 from agentops import ActionEvent, ErrorEvent
 
 
-@pytest.fixture(autouse=True, scope="function")
-def mock_req():
-    with requests_mock.Mocker() as m:
-        url = "https://api.agentops.ai"
-        m.post(url + "/v2/create_events", json={"status": "ok"})
-        m.post(url + "/v2/create_session", json={"status": "success", "jwt": "some_jwt"})
-        m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", json={"status": "ok"})
-        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
-
-        m.post(url + "/v2/reauthorize_jwt", json={"status": "success", "jwt": "some_jwt"})
-        yield m
-
-
 class TestEvents:
     def setup_method(self):
         self.api_key = "11111111-1111-4111-8111-111111111111"

tests/unit/test_pre_init.py

@@ -1,29 +1,13 @@
+import contextlib
+import time
+from datetime import datetime
+
 import pytest
 import requests_mock
-import time
+
 import agentops
 from agentops import record_action, track_agent
-from datetime import datetime
 from agentops.singleton import clear_singletons
-import contextlib
-
-jwts = ["some_jwt", "some_jwt2", "some_jwt3"]
-
-
-@contextlib.contextmanager
-@pytest.fixture(autouse=True, scope="function")
-def mock_req():
-    with requests_mock.Mocker() as m:
-        url = "https://api.agentops.ai"
-        m.post(url + "/v2/create_agent", json={"status": "success"})
-        m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/create_session", json={"status": "success", "jwt": "some_jwt"})
-        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
-
-        m.post(url + "/v2/create_events", json={"status": "ok"})
-        m.post(url + "/v2/developer_errors", json={"status": "ok"})
-
-        yield m
 
 
 @track_agent(name="TestAgent")

tests/unit/test_record_action.py

@@ -1,39 +1,10 @@
-import contextlib
 import time
 from datetime import datetime
 
 import pytest
-import requests_mock
 
 import agentops
 from agentops import record_action
-from agentops.singleton import clear_singletons
-
-jwts = ["some_jwt", "some_jwt2", "some_jwt3"]
-
-
-@contextlib.contextmanager
-@pytest.fixture(autouse=True, scope="function")
-def mock_req():
-    with requests_mock.Mocker() as m:
-        url = "https://api.agentops.ai"
-        m.post(url + "/v2/create_events", json={"status": "ok"})
-
-        # Use iter to create an iterator that can return the jwt values
-        jwt_tokens = iter(jwts)
-
-        # Use an inner function to change the response for each request
-        def create_session_response(request, context):
-            context.status_code = 200
-            return {"status": "success", "jwt": next(jwt_tokens)}
-
-        m.post(url + "/v2/create_session", json=create_session_response)
-        m.post(url + "/v2/create_events", json={"status": "ok"})
-        m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", json={"status": "ok"})
-        m.post("https://pypi.org/pypi/agentops/json", status_code=404)
-
-        yield m
 
 
 class TestRecordAction:
@@ -165,14 +136,20 @@ def add_three(x, y, z=3):
 
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.last_request.headers["Authorization"] == "Bearer some_jwt2"
+        assert (
+            mock_req.last_request.headers["Authorization"]
+            == f"Bearer {mock_req.session_jwts[str(session_2.session_id)]}"
+        )
         assert request_json["events"][0]["action_type"] == self.event_type
         assert request_json["events"][0]["params"] == {"x": 1, "y": 2, "z": 3}
         assert request_json["events"][0]["returns"] == 6
 
         second_last_request_json = mock_req.request_history[-2].json()
         assert mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.request_history[-2].headers["Authorization"] == "Bearer some_jwt"
+        assert (
+            mock_req.request_history[-2].headers["Authorization"]
+            == f"Bearer {mock_req.session_jwts[str(session_1.session_id)]}"
+        )
         assert second_last_request_json["events"][0]["action_type"] == self.event_type
         assert second_last_request_json["events"][0]["params"] == {
             "x": 1,
@@ -208,14 +185,20 @@ async def async_add(x, y):
 
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.last_request.headers["Authorization"] == "Bearer some_jwt2"
+        assert (
+            mock_req.last_request.headers["Authorization"]
+            == f"Bearer {mock_req.session_jwts[str(session_2.session_id)]}"
+        )
         assert request_json["events"][0]["action_type"] == self.event_type
         assert request_json["events"][0]["params"] == {"x": 1, "y": 2}
         assert request_json["events"][0]["returns"] == 3
 
         second_last_request_json = mock_req.request_history[-2].json()
         assert mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.request_history[-2].headers["Authorization"] == "Bearer some_jwt"
+        assert (
+            mock_req.request_history[-2].headers["Authorization"]
+            == f"Bearer {mock_req.session_jwts[str(session_1.session_id)]}"
+        )
         assert second_last_request_json["events"][0]["action_type"] == self.event_type
         assert second_last_request_json["events"][0]["params"] == {
             "x": 1,
@@ -226,7 +209,7 @@ async def async_add(x, y):
         session_1.end_session(end_state="Success")
         session_2.end_session(end_state="Success")
 
-    def test_require_session_if_multiple(self):
+    def test_require_session_if_multiple(self, mock_req):
         session_1 = agentops.start_session()
         session_2 = agentops.start_session()
 

tests/unit/test_record_tool.py

@@ -11,26 +11,6 @@
 jwts = ["some_jwt", "some_jwt2", "some_jwt3"]
 
 
-@contextlib.contextmanager
-@pytest.fixture(autouse=True)
-def mock_req():
-    with requests_mock.Mocker() as m:
-        url = "https://api.agentops.ai"
-        m.post(url + "/v2/create_events", json={"status": "ok"})
-
-        # Use iter to create an iterator that can return the jwt values
-        jwt_tokens = iter(jwts)
-
-        # Use an inner function to change the response for each request
-        def create_session_response(request, context):
-            context.status_code = 200
-            return {"status": "success", "jwt": next(jwt_tokens)}
-
-        m.post(url + "/v2/create_session", json=create_session_response)
-        m.post(url + "/v2/update_session", json={"status": "success", "token_cost": 5})
-        m.post(url + "/v2/developer_errors", json={"status": "ok"})
-
-        yield m
 
 
 class TestRecordTool:
@@ -158,14 +138,14 @@ def add_three(x, y, z=3):
 
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.last_request.headers["Authorization"] == "Bearer some_jwt2"
+        assert mock_req.last_request.headers["Authorization"] == f"Bearer {mock_req.session_jwts[str(session_2.session_id)]}"
         assert request_json["events"][0]["name"] == self.tool_name
         assert request_json["events"][0]["params"] == {"x": 1, "y": 2, "z": 3}
         assert request_json["events"][0]["returns"] == 6
 
         second_last_request_json = mock_req.request_history[-2].json()
         assert mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.request_history[-2].headers["Authorization"] == "Bearer some_jwt"
+        assert mock_req.request_history[-2].headers["Authorization"] == f"Bearer {mock_req.session_jwts[str(session_1.session_id)]}"
         assert second_last_request_json["events"][0]["name"] == self.tool_name
         assert second_last_request_json["events"][0]["params"] == {
             "x": 1,
@@ -201,14 +181,14 @@ async def async_add(x, y):
 
         request_json = mock_req.last_request.json()
         assert mock_req.last_request.headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.last_request.headers["Authorization"] == "Bearer some_jwt2"
+        assert mock_req.last_request.headers["Authorization"] == f"Bearer {mock_req.session_jwts[str(session_2.session_id)]}"
         assert request_json["events"][0]["name"] == self.tool_name
         assert request_json["events"][0]["params"] == {"x": 1, "y": 2}
         assert request_json["events"][0]["returns"] == 3
 
         second_last_request_json = mock_req.request_history[-2].json()
         assert mock_req.request_history[-2].headers["X-Agentops-Api-Key"] == self.api_key
-        assert mock_req.request_history[-2].headers["Authorization"] == "Bearer some_jwt"
+        assert mock_req.request_history[-2].headers["Authorization"] == f"Bearer {mock_req.session_jwts[str(session_1.session_id)]}"
         assert second_last_request_json["events"][0]["name"] == self.tool_name
         assert second_last_request_json["events"][0]["params"] == {
             "x": 1,
@@ -219,7 +199,7 @@ async def async_add(x, y):
         session_1.end_session(end_state="Success")
         session_2.end_session(end_state="Success")
 
-    def test_require_session_if_multiple(self):
+    def test_require_session_if_multiple(self, mock_req):
         session_1 = agentops.start_session()
         session_2 = agentops.start_session()