Testcases

Author: UtkarshMishra-Microsoft

src/backend/tests/context/test_cosmos_memory.py

@@ -13,11 +13,13 @@
 os.environ["AZURE_OPENAI_API_VERSION"] = "2023-01-01"
 os.environ["AZURE_OPENAI_ENDPOINT"] = "https://mock-openai-endpoint"
 
+
 async def async_iterable(mock_items):
     """Helper to create an async iterable."""
     for item in mock_items:
         yield item
 
+
 @pytest.fixture(autouse=True)
 def mock_env_variables(monkeypatch):
     """Mock all required environment variables."""
@@ -33,12 +35,6 @@ def mock_env_variables(monkeypatch):
     for key, value in env_vars.items():
         monkeypatch.setenv(key, value)
 
-@pytest.fixture(autouse=True)
-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():
@@ -48,6 +44,7 @@ def mock_cosmos_client():
     mock_client.create_container_if_not_exists.return_value = mock_container
     return mock_client, mock_container
 
+
 @pytest.fixture
 def mock_config(mock_cosmos_client):
     """Fixture to patch Config with mock Cosmos DB client."""
@@ -57,21 +54,20 @@ def mock_config(mock_cosmos_client):
     ), patch("src.backend.config.Config.COSMOSDB_CONTAINER", "mock-container"):
         yield
 
+
 @pytest.mark.asyncio
 async def test_initialize(mock_config, mock_cosmos_client):
     """Test if the Cosmos DB container is initialized correctly."""
     mock_client, mock_container = mock_cosmos_client
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        await context.initialize()
-        mock_client.create_container_if_not_exists.assert_called_once_with(
-            id="mock-container", partition_key=PartitionKey(path="/session_id")
-        )
-        assert context._container == mock_container
-    finally:
-        await context.close()
+    await context.initialize()
+    mock_client.create_container_if_not_exists.assert_called_once_with(
+        id="mock-container", partition_key=PartitionKey(path="/session_id")
+    )
+    assert context._container == mock_container
+
 
 @pytest.mark.asyncio
 async def test_add_item(mock_config, mock_cosmos_client):
@@ -83,14 +79,13 @@ async def test_add_item(mock_config, mock_cosmos_client):
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        await context.initialize()
-        await context.add_item(mock_item)
-        mock_container.create_item.assert_called_once_with(
-            body={"id": "test-item", "data": "test-data"}
-        )
-    finally:
-        await context.close()
+    await context.initialize()
+    await context.add_item(mock_item)
+
+    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):
@@ -102,14 +97,13 @@ async def test_update_item(mock_config, mock_cosmos_client):
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        await context.initialize()
-        await context.update_item(mock_item)
-        mock_container.upsert_item.assert_called_once_with(
-            body={"id": "test-item", "data": "updated-data"}
-        )
-    finally:
-        await context.close()
+    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):
@@ -124,72 +118,128 @@ async def test_get_item_by_id(mock_config, mock_cosmos_client):
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        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"
-        )
-    finally:
-        await context.close()
+    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
+
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        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"
-        )
-    finally:
-        await context.close()
+    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_get_data_by_invalid_type(mock_config, mock_cosmos_client):
-    """Test querying data with an invalid type."""
+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"}
+
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        result = await context.get_data_by_type("invalid_type")
-        assert result == []  # Expect empty result for invalid type
-    finally:
-        await context.close()
+    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_get_plan_by_invalid_session(mock_config, mock_cosmos_client):
-    """Test retrieving a plan with an invalid session ID."""
+async def test_update_plan(mock_config, mock_cosmos_client):
+    """Test updating a plan in Cosmos DB."""
     _, mock_container = mock_cosmos_client
-    mock_container.query_items.return_value = async_iterable([])  # No results
+    mock_plan = MagicMock()
+    mock_plan.model_dump.return_value = {"id": "plan1", "data": "updated-plan-data"}
 
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        await context.initialize()
-        result = await context.get_plan_by_session("invalid_session")
-        assert result is None
-    finally:
-        await context.close()
+    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_close_without_initialization(mock_config):
-    """Test close method without prior initialization."""
+async def test_add_session(mock_config, mock_cosmos_client):
+    """Test adding a session to Cosmos DB."""
+    _, mock_container = mock_cosmos_client
+    mock_session = MagicMock()
+    mock_session.model_dump.return_value = {"id": "session1", "data": "session-data"}
+
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        await context.close()
-    except Exception as e:
-        pytest.fail(f"Unexpected exception during close: {e}")
+    await context.initialize()
+    await context.add_session(mock_session)
+
+    mock_container.create_item.assert_called_once_with(
+        body={"id": "session1", "data": "session-data"}
+    )
+
+
+@pytest.mark.asyncio
+async def test_initialize_event(mock_config, mock_cosmos_client):
+    """Test the initialization event is set."""
+    _, _ = mock_cosmos_client
+    context = CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    )
+    assert not context._initialized.is_set()
+    await context.initialize()
+    assert context._initialized.is_set()
+
+
+@pytest.mark.asyncio
+async def test_get_data_by_invalid_type(mock_config, mock_cosmos_client):
+    """Test querying data with an invalid type."""
+    _, _ = mock_cosmos_client
+    context = CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    )
+
+    result = await context.get_data_by_type("invalid_type")
+
+    assert result == []  # Expect empty result for invalid type
+
+
+@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
+
+    context = CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    )
+    await context.initialize()
+    result = await context.get_plan_by_session("invalid_session")
+
+    assert result is None
+
 
 @pytest.mark.asyncio
 async def test_delete_item_error_handling(mock_config, mock_cosmos_client):
@@ -200,8 +250,17 @@ async def test_delete_item_error_handling(mock_config, mock_cosmos_client):
     context = CosmosBufferedChatCompletionContext(
         session_id="test_session", user_id="test_user"
     )
-    try:
-        await context.initialize()
-        await context.delete_item("test-item", "test-partition")
-    finally:
-        await context.close()
+    await context.initialize()
+    await context.delete_item(
+        "test-item", "test-partition"
+    )  # Expect no exception to propagate
+
+
+@pytest.mark.asyncio
+async def test_close_without_initialization(mock_config, mock_cosmos_client):
+    """Test close method without prior initialization."""
+    context = CosmosBufferedChatCompletionContext(
+        session_id="test_session", user_id="test_user"
+    )
+    # Expect no exceptions when closing uninitialized context
+    await context.close()