mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Added tests
This commit is contained in:
Dmytro Struk
@@ -0,0 +1,415 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import os
|
||||
from typing import Annotated
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import Field
|
||||
|
||||
from agent_framework import (
|
||||
ChatClient,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.exceptions import ServiceInitializationError
|
||||
from agent_framework.openai import OpenAIAssistantsClient
|
||||
|
||||
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
|
||||
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
|
||||
or os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
|
||||
reason="No real OPENAI_API_KEY provided; skipping integration tests."
|
||||
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
|
||||
else "Integration tests are disabled.",
|
||||
)
|
||||
|
||||
|
||||
def create_test_openai_assistants_client(
|
||||
mock_async_openai: MagicMock,
|
||||
ai_model_id: str | None = None,
|
||||
assistant_id: str | None = None,
|
||||
assistant_name: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
should_delete_assistant: bool = False,
|
||||
) -> OpenAIAssistantsClient:
|
||||
"""Helper function to create OpenAIAssistantsClient instances for testing, bypassing Pydantic validation."""
|
||||
return OpenAIAssistantsClient.model_construct(
|
||||
ai_model_id=ai_model_id or "gpt-4",
|
||||
assistant_id=assistant_id,
|
||||
assistant_name=assistant_name,
|
||||
thread_id=thread_id,
|
||||
api_key="test-api-key",
|
||||
org_id="test-org-id",
|
||||
client=mock_async_openai,
|
||||
_should_delete_assistant=should_delete_assistant,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_async_openai() -> MagicMock:
|
||||
"""Mock AsyncOpenAI client."""
|
||||
mock_client = MagicMock()
|
||||
|
||||
# Mock beta.assistants
|
||||
mock_client.beta.assistants.create = AsyncMock(return_value=MagicMock(id="test-assistant-id"))
|
||||
mock_client.beta.assistants.delete = AsyncMock()
|
||||
|
||||
# Mock beta.threads
|
||||
mock_client.beta.threads.create = AsyncMock(return_value=MagicMock(id="test-thread-id"))
|
||||
mock_client.beta.threads.delete = AsyncMock()
|
||||
|
||||
# Mock beta.threads.runs
|
||||
mock_client.beta.threads.runs.create = AsyncMock(return_value=MagicMock(id="test-run-id"))
|
||||
mock_client.beta.threads.runs.retrieve = AsyncMock()
|
||||
mock_client.beta.threads.runs.submit_tool_outputs = AsyncMock()
|
||||
|
||||
# Mock beta.threads.messages
|
||||
mock_client.beta.threads.messages.create = AsyncMock()
|
||||
mock_client.beta.threads.messages.list = AsyncMock(return_value=MagicMock(data=[]))
|
||||
|
||||
return mock_client
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_with_client(mock_async_openai: MagicMock) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with existing client."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, ai_model_id="gpt-4", assistant_id="existing-assistant-id", thread_id="test-thread-id"
|
||||
)
|
||||
|
||||
assert chat_client.client is mock_async_openai
|
||||
assert chat_client.ai_model_id == "gpt-4"
|
||||
assert chat_client.assistant_id == "existing-assistant-id"
|
||||
assert chat_client.thread_id == "test-thread-id"
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
assert isinstance(chat_client, ChatClient)
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_auto_create_client(
|
||||
openai_unit_test_env: dict[str, str],
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with auto-created client."""
|
||||
chat_client = OpenAIAssistantsClient.model_construct(
|
||||
ai_model_id=openai_unit_test_env["OPENAI_CHAT_MODEL_ID"],
|
||||
assistant_id=None,
|
||||
assistant_name="TestAssistant",
|
||||
thread_id=None,
|
||||
api_key=openai_unit_test_env["OPENAI_API_KEY"],
|
||||
org_id=openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
client=mock_async_openai,
|
||||
_should_delete_assistant=False,
|
||||
)
|
||||
|
||||
assert chat_client.client is mock_async_openai
|
||||
assert chat_client.ai_model_id == openai_unit_test_env["OPENAI_CHAT_MODEL_ID"]
|
||||
assert chat_client.assistant_id is None
|
||||
assert chat_client.assistant_name == "TestAssistant"
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_validation_fail() -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with validation failure."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
# Force failure by providing invalid model ID type - this should cause validation to fail
|
||||
OpenAIAssistantsClient(ai_model_id=123, api_key="valid-key") # type: ignore
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_openai_assistants_client_init_missing_model_id(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with missing model ID."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIAssistantsClient(
|
||||
api_key=openai_unit_test_env.get("OPENAI_API_KEY", "test-key"), env_file_path="nonexistent.env"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["OPENAI_API_KEY"]], indirect=True)
|
||||
def test_openai_assistants_client_init_missing_api_key(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with missing API key."""
|
||||
with pytest.raises(ServiceInitializationError):
|
||||
OpenAIAssistantsClient(ai_model_id="gpt-4", env_file_path="nonexistent.env")
|
||||
|
||||
|
||||
def test_openai_assistants_client_init_with_default_headers(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test OpenAIAssistantsClient initialization with default headers."""
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
chat_client = OpenAIAssistantsClient(
|
||||
ai_model_id="gpt-4",
|
||||
api_key=openai_unit_test_env["OPENAI_API_KEY"],
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
assert chat_client.ai_model_id == "gpt-4"
|
||||
assert isinstance(chat_client, ChatClient)
|
||||
|
||||
# Assert that the default header we added is present in the client's default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in chat_client.client.default_headers
|
||||
assert chat_client.client.default_headers[key] == value
|
||||
|
||||
|
||||
async def test_openai_assistants_client_get_assistant_id_or_create_existing_assistant(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_assistant_id_or_create when assistant_id is already provided."""
|
||||
chat_client = create_test_openai_assistants_client(mock_async_openai, assistant_id="existing-assistant-id")
|
||||
|
||||
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
|
||||
|
||||
assert assistant_id == "existing-assistant-id"
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
mock_async_openai.beta.assistants.create.assert_not_called()
|
||||
|
||||
|
||||
async def test_openai_assistants_client_get_assistant_id_or_create_create_new(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _get_assistant_id_or_create when creating a new assistant."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, ai_model_id="gpt-4", assistant_name="TestAssistant"
|
||||
)
|
||||
|
||||
assistant_id = await chat_client._get_assistant_id_or_create() # type: ignore
|
||||
|
||||
assert assistant_id == "test-assistant-id"
|
||||
assert chat_client._should_delete_assistant # type: ignore
|
||||
mock_async_openai.beta.assistants.create.assert_called_once()
|
||||
|
||||
|
||||
async def test_openai_assistants_client_cleanup_assistant_if_needed_should_delete(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _cleanup_assistant_if_needed when assistant should be deleted."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
|
||||
)
|
||||
|
||||
await chat_client._cleanup_assistant_if_needed() # type: ignore
|
||||
|
||||
# Verify assistant deletion was called
|
||||
mock_async_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
|
||||
|
||||
async def test_openai_assistants_client_cleanup_assistant_if_needed_should_not_delete(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _cleanup_assistant_if_needed when assistant should not be deleted."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-keep", should_delete_assistant=False
|
||||
)
|
||||
|
||||
await chat_client._cleanup_assistant_if_needed() # type: ignore
|
||||
|
||||
# Verify assistant deletion was not called
|
||||
mock_async_openai.beta.assistants.delete.assert_not_called()
|
||||
assert not chat_client._should_delete_assistant # type: ignore
|
||||
|
||||
|
||||
async def test_openai_assistants_client_cleanup_assistant_if_needed_exception_handling(
|
||||
mock_async_openai: MagicMock,
|
||||
) -> None:
|
||||
"""Test _cleanup_assistant_if_needed propagates exceptions (it doesn't handle them)."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
|
||||
)
|
||||
mock_async_openai.beta.assistants.delete.side_effect = Exception("Deletion failed")
|
||||
|
||||
with pytest.raises(Exception, match="Deletion failed"):
|
||||
await chat_client._cleanup_assistant_if_needed() # type: ignore
|
||||
|
||||
|
||||
async def test_openai_assistants_client_aclose(mock_async_openai: MagicMock) -> None:
|
||||
"""Test close method calls cleanup."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
|
||||
)
|
||||
|
||||
await chat_client.close()
|
||||
|
||||
# Verify assistant deletion was called
|
||||
mock_async_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
|
||||
|
||||
|
||||
async def test_openai_assistants_client_async_context_manager(mock_async_openai: MagicMock) -> None:
|
||||
"""Test async context manager functionality."""
|
||||
chat_client = create_test_openai_assistants_client(
|
||||
mock_async_openai, assistant_id="assistant-to-delete", should_delete_assistant=True
|
||||
)
|
||||
|
||||
# Test context manager
|
||||
async with chat_client:
|
||||
pass # Just test that we can enter and exit
|
||||
|
||||
# Verify cleanup was called on exit
|
||||
mock_async_openai.beta.assistants.delete.assert_called_once_with("assistant-to-delete")
|
||||
|
||||
|
||||
def test_openai_assistants_client_serialize(openai_unit_test_env: dict[str, str]) -> None:
|
||||
"""Test serialization of OpenAIAssistantsClient."""
|
||||
default_headers = {"X-Unit-Test": "test-guid"}
|
||||
|
||||
# Test basic initialization and to_dict
|
||||
chat_client = OpenAIAssistantsClient(
|
||||
ai_model_id="gpt-4",
|
||||
assistant_id="test-assistant-id",
|
||||
assistant_name="TestAssistant",
|
||||
thread_id="test-thread-id",
|
||||
api_key=openai_unit_test_env["OPENAI_API_KEY"],
|
||||
org_id=openai_unit_test_env["OPENAI_ORG_ID"],
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
dumped_settings = chat_client.to_dict()
|
||||
|
||||
assert dumped_settings["ai_model_id"] == "gpt-4"
|
||||
assert dumped_settings["assistant_id"] == "test-assistant-id"
|
||||
assert dumped_settings["assistant_name"] == "TestAssistant"
|
||||
assert dumped_settings["thread_id"] == "test-thread-id"
|
||||
assert dumped_settings["api_key"] == openai_unit_test_env["OPENAI_API_KEY"]
|
||||
assert dumped_settings["org_id"] == openai_unit_test_env["OPENAI_ORG_ID"]
|
||||
|
||||
# Assert that the default header we added is present in the dumped_settings default headers
|
||||
for key, value in default_headers.items():
|
||||
assert key in dumped_settings["default_headers"]
|
||||
assert dumped_settings["default_headers"][key] == value
|
||||
# Assert that the 'User-Agent' header is not present in the dumped_settings default headers
|
||||
assert "User-Agent" not in dumped_settings["default_headers"]
|
||||
|
||||
|
||||
def get_weather(
|
||||
location: Annotated[str, Field(description="The location to get the weather for.")],
|
||||
) -> str:
|
||||
"""Get the weather for a given location."""
|
||||
return f"The weather in {location} is sunny with a high of 25°C."
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_get_response() -> None:
|
||||
"""Test OpenAI Assistants Client response."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="The weather in Seattle is currently sunny with a high of 25°C. "
|
||||
"It's a beautiful day for outdoor activities.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_assistants_client.get_response(messages=messages)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert any(word in response.text.lower() for word in ["sunny", "25", "weather", "seattle"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_get_response_tools() -> None:
|
||||
"""Test OpenAI Assistants Client response with tools."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_assistants_client.get_response(
|
||||
messages=messages,
|
||||
tools=[get_weather],
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert any(word in response.text.lower() for word in ["sunny", "25", "weather"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_streaming() -> None:
|
||||
"""Test OpenAI Assistants Client streaming response."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(
|
||||
ChatMessage(
|
||||
role="user",
|
||||
text="The weather in Seattle is currently sunny with a high of 25°C. "
|
||||
"It's a beautiful day for outdoor activities.",
|
||||
)
|
||||
)
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like today?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = openai_assistants_client.get_streaming_response(messages=messages)
|
||||
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert any(word in full_message.lower() for word in ["sunny", "25", "weather", "seattle"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_streaming_tools() -> None:
|
||||
"""Test OpenAI Assistants Client streaming response with tools."""
|
||||
async with OpenAIAssistantsClient() as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
|
||||
messages: list[ChatMessage] = []
|
||||
messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?"))
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = openai_assistants_client.get_streaming_response(
|
||||
messages=messages,
|
||||
tools=[get_weather],
|
||||
tool_choice="auto",
|
||||
)
|
||||
full_message: str = ""
|
||||
async for chunk in response:
|
||||
assert chunk is not None
|
||||
assert isinstance(chunk, ChatResponseUpdate)
|
||||
for content in chunk.contents:
|
||||
if isinstance(content, TextContent) and content.text:
|
||||
full_message += content.text
|
||||
|
||||
assert any(word in full_message.lower() for word in ["sunny", "25", "weather"])
|
||||
|
||||
|
||||
@skip_if_openai_integration_tests_disabled
|
||||
async def test_openai_assistants_client_with_existing_assistant() -> None:
|
||||
"""Test OpenAI Assistants Client with existing assistant ID."""
|
||||
# First create an assistant to use in the test
|
||||
async with OpenAIAssistantsClient() as temp_client:
|
||||
# Get the assistant ID by triggering assistant creation
|
||||
messages = [ChatMessage(role="user", text="Hello")]
|
||||
await temp_client.get_response(messages=messages)
|
||||
assistant_id = temp_client.assistant_id
|
||||
|
||||
# Now test using the existing assistant
|
||||
async with OpenAIAssistantsClient(
|
||||
ai_model_id="gpt-4o-mini", assistant_id=assistant_id
|
||||
) as openai_assistants_client:
|
||||
assert isinstance(openai_assistants_client, ChatClient)
|
||||
assert openai_assistants_client.assistant_id == assistant_id
|
||||
|
||||
messages = [ChatMessage(role="user", text="What can you do?")]
|
||||
|
||||
# Test that the client can be used to get a response
|
||||
response = await openai_assistants_client.get_response(messages=messages)
|
||||
|
||||
assert response is not None
|
||||
assert isinstance(response, ChatResponse)
|
||||
assert len(response.text) > 0
|
||||
Reference in New Issue
Block a user