feat library tests added

Author: ruFFaa

agentvault_server_sdk/tests/test_fastapi_integration.py

@@ -193,13 +193,25 @@ def make_rpc_request(
     client: TestClient,
     method: str,
     params: Optional[Dict[str, Any]] = None,
-    req_id: Union[str, int] = 1
+    req_id: Union[str, int] = 1,
+    # --- ADDED: Allow sending raw content ---
+    raw_content: Optional[Union[str, bytes]] = None,
+    json_payload: Optional[Dict[str, Any]] = None,
+    # --- END ADDED ---
 ) -> Any:
     """Helper to make JSON-RPC POST requests."""
-    payload = {"jsonrpc": "2.0", "method": method, "id": req_id}
-    if params is not None:
-        payload["params"] = params
-    response = client.post("/a2a/", json=payload)
+    # --- MODIFIED: Handle raw_content ---
+    if raw_content is not None:
+        headers = {"Content-Type": "application/json"} # Still set header
+        response = client.post("/a2a/", content=raw_content, headers=headers)
+    elif json_payload is not None:
+         response = client.post("/a2a/", json=json_payload)
+    # --- END MODIFIED ---
+    else:
+        payload = {"jsonrpc": "2.0", "method": method, "id": req_id}
+        if params is not None:
+            payload["params"] = params
+        response = client.post("/a2a/", json=payload)
     return response
 
 # --- Test Cases ---
@@ -379,3 +391,80 @@ async def test_route_tasks_sendSubscribe_generator_error(test_app_with_agent):
     assert error_data.get("error") == "stream_error"
     assert f"RuntimeError: {error_message}" in error_data.get("message", "")
     # --- END MODIFIED ---
+
+# --- ADDED: Tests for Invalid JSON-RPC Requests ---
+
+def test_invalid_json_request(test_app_with_agent):
+    """Test sending invalid JSON."""
+    _, client, _ = test_app_with_agent
+    response = make_rpc_request(client, method="", raw_content=b"{invalid json")
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] is None # No ID could be parsed
+    assert resp_data["error"]["code"] == JSONRPC_PARSE_ERROR
+
+def test_non_dict_request(test_app_with_agent):
+    """Test sending a JSON array instead of an object."""
+    _, client, _ = test_app_with_agent
+    response = make_rpc_request(client, method="", json_payload=[1, 2, 3])
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] is None # No ID could be parsed
+    assert resp_data["error"]["code"] == JSONRPC_INVALID_REQUEST
+    assert "Payload must be a JSON object" in resp_data["error"]["message"]
+
+def test_missing_method_request(test_app_with_agent):
+    """Test sending a request missing the 'method' field."""
+    _, client, _ = test_app_with_agent
+    payload = {"jsonrpc": "2.0", "id": "m-err"}
+    response = make_rpc_request(client, method="", json_payload=payload)
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] == "m-err"
+    assert resp_data["error"]["code"] == JSONRPC_INVALID_REQUEST
+    assert "'method' is required" in resp_data["error"]["message"]
+
+def test_invalid_jsonrpc_version(test_app_with_agent):
+    """Test sending a request with the wrong 'jsonrpc' version."""
+    _, client, _ = test_app_with_agent
+    payload = {"jsonrpc": "1.0", "method": "test", "id": "v-err"}
+    response = make_rpc_request(client, method="", json_payload=payload)
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] == "v-err"
+    assert resp_data["error"]["code"] == JSONRPC_INVALID_REQUEST
+    assert "'jsonrpc' must be '2.0'" in resp_data["error"]["message"]
+
+# --- ADDED: Tests for tasks/sendSubscribe Parameter Validation ---
+
+def test_route_tasks_sendSubscribe_missing_params(test_app_with_agent):
+    """Test tasks/sendSubscribe with missing 'params' field."""
+    mock_agent, client, store = test_app_with_agent
+    response = make_rpc_request(client, "tasks/sendSubscribe", params=None, req_id="sub-bad1")
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] == "sub-bad1"
+    assert resp_data["error"]["code"] == JSONRPC_INVALID_PARAMS
+    assert "Params must be a dictionary" in resp_data["error"]["message"]
+
+def test_route_tasks_sendSubscribe_missing_id_in_params(test_app_with_agent):
+    """Test tasks/sendSubscribe with 'params' missing the 'id' key."""
+    mock_agent, client, store = test_app_with_agent
+    response = make_rpc_request(client, "tasks/sendSubscribe", params={"other": "value"}, req_id="sub-bad2")
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] == "sub-bad2"
+    assert resp_data["error"]["code"] == JSONRPC_INVALID_PARAMS
+    assert "'id' parameter is required" in resp_data["error"]["message"]
+
+def test_route_tasks_sendSubscribe_invalid_id_type(test_app_with_agent):
+    """Test tasks/sendSubscribe with 'id' parameter of wrong type."""
+    mock_agent, client, store = test_app_with_agent
+    response = make_rpc_request(client, "tasks/sendSubscribe", params={"id": 12345}, req_id="sub-bad3")
+    assert response.status_code == status.HTTP_200_OK # JSON-RPC error
+    resp_data = response.json()
+    assert resp_data["id"] == "sub-bad3"
+    assert resp_data["error"]["code"] == JSONRPC_INVALID_PARAMS
+    assert "'id' parameter is required" in resp_data["error"]["message"] # Error message might be generic
+
+# --- END ADDED ---

agentvault_server_sdk/tests/test_state.py

@@ -4,7 +4,10 @@
 from freezegun import freeze_time
 from typing import Union, Any, List
 # --- ADDED: Import patch ---
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock, AsyncMock # Added AsyncMock
+# --- END ADDED ---
+# --- ADDED: Import logging ---
+import logging
 # --- END ADDED ---
 
 
@@ -339,9 +342,119 @@ async def test_update_task_state_triggers_notify(task_store: InMemoryTaskStore):
     # Patch the notify method to check it's called
     with patch.object(task_store, 'notify_status_update', wraps=task_store.notify_status_update) as mock_notify:
         await task_store.update_task_state(task_id, TaskState.WORKING)
-        mock_notify.assert_awaited_once_with(task_id, TaskState.WORKING)
+        # --- MODIFIED: Correct assertion ---
+        mock_notify.assert_awaited_once_with(task_id, TaskState.WORKING) # Removed None
+        # --- END MODIFIED ---
 
     # Check the queue still received the event via the wrapped call
     event = await asyncio.wait_for(q1.get(), timeout=0.1)
     assert isinstance(event, TaskStatusUpdateEvent)
     assert event.state == TaskState.WORKING
+
+# --- ADDED: Tests for Edge Cases and Errors ---
+
+@pytest.mark.asyncio
+async def test_create_task_already_exists(task_store: InMemoryTaskStore, caplog):
+    """Test creating a task that already exists."""
+    task_id = "existing-task"
+    await task_store.create_task(task_id) # Create it first
+    initial_context = task_store._tasks[task_id]
+
+    with caplog.at_level(logging.WARNING):
+        new_context = await task_store.create_task(task_id) # Try creating again
+
+    assert new_context is initial_context # Should return the existing one
+    assert f"Task '{task_id}' already exists" in caplog.text
+
+@pytest.mark.asyncio
+async def test_update_task_state_not_found(task_store: InMemoryTaskStore, caplog):
+    """Test updating state for a non-existent task."""
+    task_id = "non-existent-task"
+    with caplog.at_level(logging.WARNING):
+        context = await task_store.update_task_state(task_id, TaskState.WORKING)
+
+    assert context is None
+    assert f"Task '{task_id}' not found for state update" in caplog.text
+
+@pytest.mark.asyncio
+async def test_delete_task_not_found(task_store: InMemoryTaskStore, caplog):
+    """Test deleting a non-existent task."""
+    task_id = "non-existent-task"
+    with caplog.at_level(logging.WARNING):
+        result = await task_store.delete_task(task_id)
+
+    assert result is False
+    assert f"Task '{task_id}' not found for deletion" in caplog.text
+
+@pytest.mark.asyncio
+async def test_remove_listener_task_not_found(task_store: InMemoryTaskStore, caplog):
+    """Test removing a listener from a non-existent task."""
+    task_id = "non-existent-task"
+    q1 = asyncio.Queue()
+    with caplog.at_level(logging.WARNING):
+        await task_store.remove_listener(task_id, q1)
+    assert f"Attempted to remove listener from non-existent task '{task_id}'" in caplog.text
+
+@pytest.mark.asyncio
+async def test_remove_listener_queue_not_found(task_store: InMemoryTaskStore, caplog):
+    """Test removing a listener queue that wasn't added."""
+    task_id = "listener-task-2"
+    q1 = asyncio.Queue()
+    q2 = asyncio.Queue()
+    await task_store.add_listener(task_id, q1) # Add q1
+
+    with caplog.at_level(logging.WARNING):
+        await task_store.remove_listener(task_id, q2) # Try removing q2
+
+    assert f"Attempted to remove a listener queue not present for task '{task_id}'" in caplog.text
+    assert await task_store.get_listeners(task_id) == [q1] # q1 should still be there
+
+@pytestmark_notify
+@pytest.mark.asyncio
+async def test_notify_listeners_queue_put_error(task_store: InMemoryTaskStore, caplog):
+    """Test error handling when putting event onto a listener queue fails."""
+    task_id = "notify-queue-error"
+    await task_store.create_task(task_id)
+    q_ok = asyncio.Queue()
+    q_bad = MagicMock(spec=asyncio.Queue)
+    q_bad.put = AsyncMock(side_effect=asyncio.QueueFull("Mock Queue Full")) # Simulate error
+
+    await task_store.add_listener(task_id, q_ok)
+    await task_store.add_listener(task_id, q_bad)
+
+    with caplog.at_level(logging.ERROR):
+        await task_store.notify_status_update(task_id, TaskState.FAILED, "Notify Error Test")
+
+    # Check the good queue received the event
+    event_ok = await asyncio.wait_for(q_ok.get(), timeout=0.1)
+    assert isinstance(event_ok, TaskStatusUpdateEvent)
+    assert event_ok.state == TaskState.FAILED
+
+    # Check the bad queue put was attempted and error logged
+    q_bad.put.assert_awaited_once()
+    assert f"Failed to put event onto listener queue 1 for task '{task_id}'" in caplog.text
+    assert "Mock Queue Full" in caplog.text
+
+@pytestmark_notify
+@pytest.mark.asyncio
+@patch("agentvault_server_sdk.state.TaskStatusUpdateEvent") # Patch the model class
+async def test_notify_status_update_event_creation_error(mock_event_cls, task_store: InMemoryTaskStore, caplog):
+    """Test error handling if TaskStatusUpdateEvent creation fails."""
+    task_id = "notify-create-error"
+    await task_store.create_task(task_id)
+    q1 = asyncio.Queue()
+    await task_store.add_listener(task_id, q1)
+
+    error_message = "Mock Pydantic validation error"
+    mock_event_cls.side_effect = ValueError(error_message) # Simulate model init error
+
+    with caplog.at_level(logging.ERROR):
+        await task_store.notify_status_update(task_id, TaskState.WORKING)
+
+    assert q1.empty() # Queue should be empty
+    assert f"Failed to create TaskStatusUpdateEvent instance for task '{task_id}'" in caplog.text
+    assert error_message in caplog.text
+
+# Similar tests can be added for notify_message_event and notify_artifact_event creation errors
+
+# --- END ADDED ---