Testcases

Author: UtkarshMishra-Microsoft

src/backend/tests/context/test_cosmos_memory.py

@@ -37,6 +37,14 @@ def mock_env_variables(monkeypatch):
         monkeypatch.setenv(key, value)
 
 
+@pytest.fixture
+def mock_azure_credentials():
+    """Mock Azure DefaultAzureCredential for all tests."""
+    with patch("azure.identity.aio.DefaultAzureCredential") as mock_cred:
+        mock_cred.return_value.get_token = AsyncMock(return_value={"token": "mock-token"})
+        yield
+
+
 @pytest.fixture
 def mock_cosmos_client():
     """Fixture for mocking Cosmos DB client and container."""
@@ -71,13 +79,12 @@ async def test_initialize(mock_config, mock_cosmos_client):
 
 
 @pytest.mark.asyncio
-async def test_close_without_initialization():
+async def test_close_without_initialization(mock_config):
     """Test closing the context without prior initialization."""
     async with CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
-    ) as context:
-        # Ensure close is safe even if not explicitly initialized
-        pass
+    ):
+        pass  # Expect no errors when exiting context
 
 
 @pytest.mark.asyncio
@@ -95,3 +102,116 @@ async def test_add_item(mock_config, mock_cosmos_client):
         mock_container.create_item.assert_called_once_with(
             body={"id": "test-item", "data": "test-data"}
         )
+
+
+@pytest.mark.asyncio
+async def test_update_item(mock_config, mock_cosmos_client):
+    """Test updating an item in Cosmos DB."""
+    _, mock_container = mock_cosmos_client
+    mock_item = MagicMock()
+    mock_item.model_dump.return_value = {"id": "test-item", "data": "updated-data"}
+
+    async with CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    ) as context:
+        await context.initialize()
+        await context.update_item(mock_item)
+        mock_container.upsert_item.assert_called_once_with(
+            body={"id": "test-item", "data": "updated-data"}
+        )
+
+
+@pytest.mark.asyncio
+async def test_get_item_by_id(mock_config, mock_cosmos_client):
+    """Test retrieving an item by ID from Cosmos DB."""
+    _, mock_container = mock_cosmos_client
+    mock_item = {"id": "test-item", "data": "retrieved-data"}
+    mock_container.read_item.return_value = mock_item
+
+    mock_model_class = MagicMock()
+    mock_model_class.model_validate.return_value = "validated_item"
+
+    async with CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    ) as context:
+        await context.initialize()
+        result = await context.get_item_by_id(
+            "test-item", "test-partition", mock_model_class
+        )
+
+        assert result == "validated_item"
+        mock_container.read_item.assert_called_once_with(
+            item="test-item", partition_key="test-partition"
+        )
+
+
+@pytest.mark.asyncio
+async def test_delete_item(mock_config, mock_cosmos_client):
+    """Test deleting an item from Cosmos DB."""
+    _, mock_container = mock_cosmos_client
+
+    async with CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    ) as context:
+        await context.initialize()
+        await context.delete_item("test-item", "test-partition")
+
+        mock_container.delete_item.assert_called_once_with(
+            item="test-item", partition_key="test-partition"
+        )
+
+
+@pytest.mark.asyncio
+async def test_add_plan(mock_config, mock_cosmos_client):
+    """Test adding a plan to Cosmos DB."""
+    _, mock_container = mock_cosmos_client
+    mock_plan = MagicMock()
+    mock_plan.model_dump.return_value = {"id": "plan1", "data": "plan-data"}
+
+    async with CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    ) as context:
+        await context.initialize()
+        await context.add_plan(mock_plan)
+
+        mock_container.create_item.assert_called_once_with(
+            body={"id": "plan1", "data": "plan-data"}
+        )
+
+
+@pytest.mark.asyncio
+async def test_update_plan(mock_config, mock_cosmos_client):
+    """Test updating a plan in Cosmos DB."""
+    _, mock_container = mock_cosmos_client
+    mock_plan = MagicMock()
+    mock_plan.model_dump.return_value = {
+        "id": "plan1",
+        "data": "updated-plan-data",
+    }
+
+    async with CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    ) as context:
+        await context.initialize()
+        await context.update_plan(mock_plan)
+
+        mock_container.upsert_item.assert_called_once_with(
+            body={"id": "plan1", "data": "updated-plan-data"}
+        )
+
+
+@pytest.mark.asyncio
+async def test_get_plan_by_invalid_session(mock_config, mock_cosmos_client):
+    """Test retrieving a plan with an invalid session ID."""
+    _, mock_container = mock_cosmos_client
+    mock_container.query_items.return_value = async_iterable(
+        []
+    )  # No results for invalid session
+
+    async with CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    ) as context:
+        await context.initialize()
+        result = await context.get_plan_by_session("invalid_session")
+
+        assert result is None