From 1316fa6747ac87d77c7d51a130293c04c30c0b0e Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Wed, 30 Jul 2025 17:55:12 -0700 Subject: [PATCH] Added tests --- .../openai/test_openai_assistants_client.py | 415 ++++++++++++++++++ 1 file changed, 415 insertions(+) create mode 100644 python/packages/main/tests/openai/test_openai_assistants_client.py diff --git a/python/packages/main/tests/openai/test_openai_assistants_client.py b/python/packages/main/tests/openai/test_openai_assistants_client.py new file mode 100644 index 0000000000..10dcf588c1 --- /dev/null +++ b/python/packages/main/tests/openai/test_openai_assistants_client.py @@ -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