hexdag 0.5.0.dev1__py3-none-any.whl

hexdag_plugins/azure/tests/test_azure_keyvault_adapter.py

@@ -0,0 +1,330 @@
+"""Tests for Azure Key Vault adapter."""
+
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from hexdag_plugins.azure.azure_keyvault_adapter import AzureKeyVaultAdapter
+
+
+@pytest.fixture
+def keyvault_adapter():
+    """Create Azure Key Vault adapter for testing."""
+    return AzureKeyVaultAdapter(
+        vault_url="https://test-vault.vault.azure.net",
+        use_managed_identity=True,
+    )
+
+
+@pytest.fixture
+def keyvault_adapter_service_principal():
+    """Create Azure Key Vault adapter with service principal auth."""
+    return AzureKeyVaultAdapter(
+        vault_url="https://test-vault.vault.azure.net",
+        use_managed_identity=False,
+        tenant_id="test-tenant-id",
+        client_id="test-client-id",
+        client_secret="test-client-secret",
+    )
+
+
+@pytest.fixture
+def keyvault_adapter_no_cache():
+    """Create adapter with caching disabled."""
+    return AzureKeyVaultAdapter(
+        vault_url="https://test-vault.vault.azure.net",
+        use_managed_identity=True,
+        cache_secrets=False,
+    )
+
+
+@pytest.mark.asyncio
+async def test_adapter_initialization(keyvault_adapter):
+    """Test adapter initializes with correct parameters."""
+    assert keyvault_adapter.vault_url == "https://test-vault.vault.azure.net"
+    assert keyvault_adapter.use_managed_identity is True
+    assert keyvault_adapter.cache_secrets is True
+    assert keyvault_adapter.cache_ttl == 300
+    assert keyvault_adapter._client is None
+
+
+@pytest.mark.asyncio
+async def test_adapter_initialization_service_principal(keyvault_adapter_service_principal):
+    """Test adapter initializes with service principal."""
+    assert keyvault_adapter_service_principal.use_managed_identity is False
+    assert keyvault_adapter_service_principal.tenant_id == "test-tenant-id"
+    assert keyvault_adapter_service_principal.client_id == "test-client-id"
+
+
+@pytest.mark.asyncio
+async def test_get_secret_success(keyvault_adapter):
+    """Test successful secret retrieval."""
+    mock_secret = MagicMock()
+    mock_secret.value = "secret-value-123"
+
+    mock_client = MagicMock()
+    mock_client.get_secret.return_value = mock_secret
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        result = await keyvault_adapter.aget("MY-SECRET")
+
+        assert result == "secret-value-123"
+        mock_client.get_secret.assert_called_once_with("MY-SECRET")
+
+
+@pytest.mark.asyncio
+async def test_get_secret_from_cache(keyvault_adapter):
+    """Test secret retrieval from cache."""
+    # Pre-populate cache
+    keyvault_adapter._cache["CACHED-SECRET"] = ("cached-value", time.time())
+
+    mock_client = MagicMock()
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        result = await keyvault_adapter.aget("CACHED-SECRET")
+
+        assert result == "cached-value"
+        mock_client.get_secret.assert_not_called()  # Should not hit the API
+
+
+@pytest.mark.asyncio
+async def test_get_secret_expired_cache(keyvault_adapter):
+    """Test secret retrieval with expired cache."""
+    # Pre-populate cache with expired entry
+    keyvault_adapter._cache["EXPIRED-SECRET"] = ("old-value", time.time() - 400)
+
+    mock_secret = MagicMock()
+    mock_secret.value = "new-value"
+
+    mock_client = MagicMock()
+    mock_client.get_secret.return_value = mock_secret
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        result = await keyvault_adapter.aget("EXPIRED-SECRET")
+
+        assert result == "new-value"
+        mock_client.get_secret.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_get_secret_no_cache(keyvault_adapter_no_cache):
+    """Test secret retrieval with caching disabled."""
+    mock_secret = MagicMock()
+    mock_secret.value = "secret-value"
+
+    mock_client = MagicMock()
+    mock_client.get_secret.return_value = mock_secret
+
+    with patch.object(keyvault_adapter_no_cache, "_get_client", return_value=mock_client):
+        # First call
+        await keyvault_adapter_no_cache.aget("NO-CACHE-SECRET")
+        # Second call should also hit the API
+        await keyvault_adapter_no_cache.aget("NO-CACHE-SECRET")
+
+        assert mock_client.get_secret.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_get_secret_not_found(keyvault_adapter):
+    """Test secret retrieval raises ValueError for missing secret."""
+    mock_client = MagicMock()
+    mock_client.get_secret.side_effect = Exception("SecretNotFound: Secret not found")
+
+    with (
+        patch.object(keyvault_adapter, "_get_client", return_value=mock_client),
+        pytest.raises(ValueError, match="not found"),
+    ):
+        await keyvault_adapter.aget("MISSING-SECRET")
+
+
+@pytest.mark.asyncio
+async def test_get_secret_null_value(keyvault_adapter):
+    """Test secret retrieval raises ValueError for null value."""
+    mock_secret = MagicMock()
+    mock_secret.value = None
+
+    mock_client = MagicMock()
+    mock_client.get_secret.return_value = mock_secret
+
+    with (
+        patch.object(keyvault_adapter, "_get_client", return_value=mock_client),
+        pytest.raises(ValueError, match="has no value"),
+    ):
+        await keyvault_adapter.aget("NULL-SECRET")
+
+
+@pytest.mark.asyncio
+async def test_get_batch_success(keyvault_adapter):
+    """Test batch secret retrieval."""
+
+    def mock_get_secret(name):
+        mock = MagicMock()
+        mock.value = f"value-for-{name}"
+        return mock
+
+    mock_client = MagicMock()
+    mock_client.get_secret.side_effect = mock_get_secret
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        results = await keyvault_adapter.aget_batch(["SECRET1", "SECRET2"])
+
+        assert len(results) == 2
+        assert results["SECRET1"] == "value-for-SECRET1"
+        assert results["SECRET2"] == "value-for-SECRET2"
+
+
+@pytest.mark.asyncio
+async def test_get_batch_partial_success(keyvault_adapter):
+    """Test batch retrieval skips missing secrets."""
+
+    def mock_get_secret(name):
+        if name == "MISSING":
+            raise Exception("SecretNotFound")
+        mock = MagicMock()
+        mock.value = f"value-for-{name}"
+        return mock
+
+    mock_client = MagicMock()
+    mock_client.get_secret.side_effect = mock_get_secret
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        results = await keyvault_adapter.aget_batch(["SECRET1", "MISSING", "SECRET2"])
+
+        assert len(results) == 2
+        assert "SECRET1" in results
+        assert "MISSING" not in results
+        assert "SECRET2" in results
+
+
+@pytest.mark.asyncio
+async def test_set_secret(keyvault_adapter):
+    """Test setting a secret."""
+    mock_client = MagicMock()
+    mock_client.set_secret = MagicMock()
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        await keyvault_adapter.aset("NEW-SECRET", "secret-value")
+
+        mock_client.set_secret.assert_called_once_with("NEW-SECRET", "secret-value")
+        # Verify it's also cached
+        assert "NEW-SECRET" in keyvault_adapter._cache
+
+
+@pytest.mark.asyncio
+async def test_delete_secret(keyvault_adapter):
+    """Test deleting a secret."""
+    # Pre-populate cache
+    keyvault_adapter._cache["TO-DELETE"] = ("value", time.time())
+
+    mock_client = MagicMock()
+    mock_client.begin_delete_secret = MagicMock()
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        await keyvault_adapter.adelete("TO-DELETE")
+
+        mock_client.begin_delete_secret.assert_called_once_with("TO-DELETE")
+        assert "TO-DELETE" not in keyvault_adapter._cache
+
+
+@pytest.mark.asyncio
+async def test_list_secrets(keyvault_adapter):
+    """Test listing all secrets."""
+    mock_secret1 = MagicMock()
+    mock_secret1.name = "SECRET1"
+    mock_secret2 = MagicMock()
+    mock_secret2.name = "SECRET2"
+
+    mock_client = MagicMock()
+    mock_client.list_properties_of_secrets.return_value = [mock_secret1, mock_secret2]
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        names = await keyvault_adapter.alist()
+
+        assert names == ["SECRET1", "SECRET2"]
+
+
+@pytest.mark.asyncio
+async def test_clear_cache(keyvault_adapter):
+    """Test clearing the cache."""
+    keyvault_adapter._cache["SECRET1"] = ("value1", time.time())
+    keyvault_adapter._cache["SECRET2"] = ("value2", time.time())
+
+    keyvault_adapter.clear_cache()
+
+    assert len(keyvault_adapter._cache) == 0
+
+
+@pytest.mark.asyncio
+async def test_health_check_healthy(keyvault_adapter):
+    """Test health check returns healthy status."""
+    mock_client = MagicMock()
+    mock_client.list_properties_of_secrets.return_value = []
+
+    with patch.object(keyvault_adapter, "_get_client", return_value=mock_client):
+        status = await keyvault_adapter.ahealth_check()
+
+        assert status.status == "healthy"
+        assert status.adapter_name == "AzureKeyVault"
+        assert status.details["vault_url"] == "https://test-vault.vault.azure.net"
+        assert status.details["auth_method"] == "managed_identity"
+
+
+@pytest.mark.asyncio
+async def test_health_check_service_principal(keyvault_adapter_service_principal):
+    """Test health check shows service principal auth method."""
+    mock_client = MagicMock()
+    mock_client.list_properties_of_secrets.return_value = []
+
+    with patch.object(keyvault_adapter_service_principal, "_get_client", return_value=mock_client):
+        status = await keyvault_adapter_service_principal.ahealth_check()
+
+        assert status.details["auth_method"] == "service_principal"
+
+
+@pytest.mark.asyncio
+async def test_health_check_unhealthy(keyvault_adapter):
+    """Test health check returns unhealthy on error."""
+    with patch.object(keyvault_adapter, "_get_client", side_effect=Exception("Connection failed")):
+        status = await keyvault_adapter.ahealth_check()
+
+        assert status.status == "unhealthy"
+        assert "error" in status.details
+
+
+@pytest.mark.asyncio
+async def test_to_dict(keyvault_adapter):
+    """Test serialization excludes credentials."""
+    config = keyvault_adapter.to_dict()
+
+    assert "vault_url" in config
+    assert "cache_secrets" in config
+    assert "client_secret" not in config  # Credentials excluded
+
+
+@pytest.mark.asyncio
+async def test_service_principal_requires_all_credentials():
+    """Test service principal auth requires all credentials."""
+    adapter = AzureKeyVaultAdapter(
+        vault_url="https://test.vault.azure.net",
+        use_managed_identity=False,
+        tenant_id="tenant",
+        # Missing client_id and client_secret
+    )
+
+    with pytest.raises(ValueError, match="tenant_id, client_id, and client_secret"):
+        adapter._get_client()
+
+
+@pytest.mark.asyncio
+async def test_client_lazy_initialization(keyvault_adapter):
+    """Test client is lazily initialized."""
+    assert keyvault_adapter._client is None
+
+    # The _get_client method creates the client
+    with (
+        patch("hexdag_plugins.azure.azure_keyvault_adapter.DefaultAzureCredential"),
+        patch("hexdag_plugins.azure.azure_keyvault_adapter.SecretClient") as mock_client,
+    ):
+        keyvault_adapter._get_client()
+        mock_client.assert_called_once()

hexdag_plugins/azure/tests/test_azure_openai_adapter.py

@@ -0,0 +1,329 @@
+"""Tests for Azure OpenAI adapter."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from hexdag.core.ports.llm import Message
+
+from hexdag_plugins.azure.azure_openai_adapter import AzureOpenAIAdapter
+
+
+@pytest.fixture
+def azure_adapter():
+    """Create Azure OpenAI adapter for testing."""
+    return AzureOpenAIAdapter(
+        api_key="test-key",
+        resource_name="test-resource",
+        deployment_id="gpt-4",
+        api_version="2024-02-15-preview",
+        temperature=0.7,
+    )
+
+
+@pytest.mark.asyncio
+async def test_adapter_initialization(azure_adapter):
+    """Test adapter initializes with correct parameters."""
+    assert azure_adapter.api_key == "test-key"
+    assert azure_adapter.resource_name == "test-resource"
+    assert azure_adapter.deployment_id == "gpt-4"
+    assert azure_adapter.api_version == "2024-02-15-preview"
+    assert azure_adapter.temperature == 0.7
+    assert azure_adapter.api_base == "https://test-resource.openai.azure.com"
+
+
+@pytest.mark.asyncio
+async def test_aresponse_success(azure_adapter):
+    """Test successful response generation."""
+    # Mock OpenAI client
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "Hello from Azure!"
+
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_adapter, "_get_client", return_value=mock_client):
+        messages = [Message(role="user", content="Hello")]
+        response = await azure_adapter.aresponse(messages)
+
+        assert response == "Hello from Azure!"
+        mock_client.chat.completions.create.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_aresponse_with_tools(azure_adapter):
+    """Test response with tool calling."""
+    from hexdag.core.ports.llm import LLMResponse
+
+    # Mock tool call response
+    mock_tool_call = MagicMock()
+    mock_tool_call.id = "call_123"
+    mock_tool_call.function.name = "search"
+    mock_tool_call.function.arguments = {"query": "test"}
+
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "Let me search"
+    mock_response.choices[0].message.tool_calls = [mock_tool_call]
+    mock_response.choices[0].finish_reason = "tool_calls"
+
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_adapter, "_get_client", return_value=mock_client):
+        messages = [Message(role="user", content="Search for cats")]
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "search",
+                    "description": "Search",
+                    "parameters": {"type": "object", "properties": {}},
+                },
+            }
+        ]
+
+        response = await azure_adapter.aresponse_with_tools(messages, tools)
+
+        assert isinstance(response, LLMResponse)
+        assert response.content == "Let me search"
+        assert len(response.tool_calls) == 1
+        assert response.tool_calls[0].name == "search"
+        assert response.finish_reason == "tool_calls"
+
+
+@pytest.mark.asyncio
+async def test_aresponse_error_handling(azure_adapter):
+    """Test error handling returns None."""
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create = AsyncMock(side_effect=Exception("API error"))
+
+    with patch.object(azure_adapter, "_get_client", return_value=mock_client):
+        messages = [Message(role="user", content="Hello")]
+        response = await azure_adapter.aresponse(messages)
+
+        assert response is None
+
+
+@pytest.mark.asyncio
+async def test_health_check_healthy(azure_adapter):
+    """Test health check returns healthy status."""
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "OK"
+
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_adapter, "_get_client", return_value=mock_client):
+        status = await azure_adapter.ahealth_check()
+
+        assert status.status == "healthy"
+        assert status.adapter_name == "AzureOpenAI[gpt-4]"
+        assert status.latency_ms > 0
+        assert status.details["resource"] == "test-resource"
+        assert status.details["deployment"] == "gpt-4"
+
+
+@pytest.mark.asyncio
+async def test_health_check_unhealthy(azure_adapter):
+    """Test health check handles errors."""
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create = AsyncMock(side_effect=Exception("Connection failed"))
+
+    with patch.object(azure_adapter, "_get_client", return_value=mock_client):
+        status = await azure_adapter.ahealth_check()
+
+        assert status.status == "unhealthy"
+        assert "error" in status.details
+
+
+@pytest.mark.asyncio
+async def test_client_lazy_initialization(azure_adapter):
+    """Test OpenAI client is lazily initialized."""
+    assert azure_adapter._client is None
+
+    # Patch at module level where it's imported
+    with patch("hexdag_plugins.azure.azure_openai_adapter.AsyncAzureOpenAI") as mock_azure:
+        mock_client = MagicMock()
+        mock_azure.return_value = mock_client
+
+        client = azure_adapter._get_client()
+
+        assert client is mock_client
+        assert azure_adapter._client is mock_client
+        mock_azure.assert_called_once_with(
+            api_key="test-key",
+            api_version="2024-02-15-preview",
+            azure_endpoint="https://test-resource.openai.azure.com",
+            timeout=30.0,
+        )
+
+
+@pytest.mark.asyncio
+async def test_client_reuse():
+    """Test that client is reused across multiple calls."""
+    adapter = AzureOpenAIAdapter(
+        api_key="test-key",
+        resource_name="test-resource",
+        deployment_id="gpt-4",
+    )
+
+    # Mock AsyncAzureOpenAI creation
+    with patch("hexdag_plugins.azure.azure_openai_adapter.AsyncAzureOpenAI") as mock_azure:
+        mock_client = MagicMock()
+        mock_azure.return_value = mock_client
+
+        # First call creates client
+        client1 = adapter._get_client()
+        assert mock_azure.call_count == 1
+
+        # Second call reuses client
+        client2 = adapter._get_client()
+        assert mock_azure.call_count == 1  # Still 1, not 2
+
+        assert client1 is client2
+        assert client1 is mock_client
+
+
+@pytest.mark.asyncio
+async def test_custom_timeout(azure_adapter):
+    """Test custom timeout is respected."""
+    custom_adapter = AzureOpenAIAdapter(
+        api_key="test-key",
+        resource_name="test-resource",
+        deployment_id="gpt-4",
+        timeout=60.0,
+    )
+
+    assert custom_adapter.timeout == 60.0
+
+
+@pytest.mark.asyncio
+async def test_max_tokens_parameter(azure_adapter):
+    """Test max_tokens parameter is passed to API."""
+    azure_adapter.max_tokens = 100
+
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "Short response"
+
+    mock_client = AsyncMock()
+    mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_adapter, "_get_client", return_value=mock_client):
+        messages = [Message(role="user", content="Hello")]
+        await azure_adapter.aresponse(messages)
+
+        call_kwargs = mock_client.chat.completions.create.call_args.kwargs
+        assert call_kwargs["max_tokens"] == 100
+
+
+# ========== Embedding Tests ==========
+
+
+@pytest.fixture
+def azure_embedding_adapter():
+    """Create Azure OpenAI adapter with embedding support."""
+    return AzureOpenAIAdapter(
+        api_key="test-key",
+        resource_name="test-resource",
+        deployment_id="gpt-4",
+        embedding_deployment_id="text-embedding-3-small",
+        embedding_dimensions=1536,
+    )
+
+
+@pytest.mark.asyncio
+async def test_aembed_success(azure_embedding_adapter):
+    """Test successful embedding generation."""
+    # Mock embedding response
+    mock_embedding_data = MagicMock()
+    mock_embedding_data.embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
+    mock_embedding_data.index = 0
+
+    mock_response = MagicMock()
+    mock_response.data = [mock_embedding_data]
+
+    mock_client = AsyncMock()
+    mock_client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_embedding_adapter, "_get_client", return_value=mock_client):
+        embedding = await azure_embedding_adapter.aembed("Hello, world!")
+
+        assert embedding == [0.1, 0.2, 0.3, 0.4, 0.5]
+        mock_client.embeddings.create.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_aembed_without_deployment_id():
+    """Test embedding fails without embedding_deployment_id."""
+    adapter = AzureOpenAIAdapter(
+        api_key="test-key",
+        resource_name="test-resource",
+        deployment_id="gpt-4",
+        # No embedding_deployment_id set
+    )
+
+    with pytest.raises(ValueError, match="embedding_deployment_id must be set"):
+        await adapter.aembed("test")
+
+
+@pytest.mark.asyncio
+async def test_aembed_batch_success(azure_embedding_adapter):
+    """Test successful batch embedding generation."""
+    # Mock batch embedding response
+    mock_data_1 = MagicMock()
+    mock_data_1.embedding = [0.1, 0.2, 0.3]
+    mock_data_1.index = 0
+
+    mock_data_2 = MagicMock()
+    mock_data_2.embedding = [0.4, 0.5, 0.6]
+    mock_data_2.index = 1
+
+    mock_response = MagicMock()
+    mock_response.data = [mock_data_1, mock_data_2]
+
+    mock_client = AsyncMock()
+    mock_client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_embedding_adapter, "_get_client", return_value=mock_client):
+        embeddings = await azure_embedding_adapter.aembed_batch(["Hello", "World"])
+
+        assert len(embeddings) == 2
+        assert embeddings[0] == [0.1, 0.2, 0.3]
+        assert embeddings[1] == [0.4, 0.5, 0.6]
+
+
+@pytest.mark.asyncio
+async def test_aembed_with_dimensions(azure_embedding_adapter):
+    """Test embedding respects dimensions parameter."""
+    mock_embedding_data = MagicMock()
+    mock_embedding_data.embedding = [0.1] * 1536
+    mock_embedding_data.index = 0
+
+    mock_response = MagicMock()
+    mock_response.data = [mock_embedding_data]
+
+    mock_client = AsyncMock()
+    mock_client.embeddings.create = AsyncMock(return_value=mock_response)
+
+    with patch.object(azure_embedding_adapter, "_get_client", return_value=mock_client):
+        await azure_embedding_adapter.aembed("test")
+
+        call_kwargs = mock_client.embeddings.create.call_args.kwargs
+        assert call_kwargs["dimensions"] == 1536
+
+
+@pytest.mark.asyncio
+async def test_aembed_image_not_implemented(azure_embedding_adapter):
+    """Test image embedding raises NotImplementedError."""
+    with pytest.raises(NotImplementedError, match="Azure OpenAI does not support image embeddings"):
+        await azure_embedding_adapter.aembed_image("image.jpg")
+
+
+@pytest.mark.asyncio
+async def test_aembed_image_batch_not_implemented(azure_embedding_adapter):
+    """Test batch image embedding raises NotImplementedError."""
+    with pytest.raises(NotImplementedError, match="Azure OpenAI does not support image embeddings"):
+        await azure_embedding_adapter.aembed_image_batch(["image1.jpg", "image2.jpg"])