tests/unit/agents/test_unit_openai.py

import os
import pytest
from unittest.mock import patch, MagicMock, ANY

from surf_spot_finder.agents.openai import (
    run_openai_agent,
    run_openai_multi_agent,
)
from surf_spot_finder.tools import (
    show_final_answer,
    show_plan,
    ask_user_verification,
    search_web,
    visit_webpage,
)
from surf_spot_finder.prompts.openai import (
    SINGLE_AGENT_SYSTEM_PROMPT,
    MULTI_AGENT_SYSTEM_PROMPT,
)


def test_run_openai_agent_default():
    mock_agent = MagicMock()

    with (
        patch("surf_spot_finder.agents.openai.Agent", mock_agent),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
    ):
        run_openai_agent("gpt-4o", "Test prompt")
        mock_agent.assert_called_once_with(
            model="gpt-4o",
            instructions=SINGLE_AGENT_SYSTEM_PROMPT,
            name="surf-spot-finder",
            tools=ANY,
        )


def test_run_openai_agent_base_url_and_api_key_var():
    async_openai_mock = MagicMock()
    openai_chat_completions_model = MagicMock()
    with (
        patch("surf_spot_finder.agents.openai.Agent", MagicMock()),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
        patch("surf_spot_finder.agents.openai.AsyncOpenAI", async_openai_mock),
        patch(
            "surf_spot_finder.agents.openai.OpenAIChatCompletionsModel",
            openai_chat_completions_model,
        ),
        patch.dict(os.environ, {"TEST_API_KEY": "test-key-12345"}),
    ):
        run_openai_agent(
            "gpt-4o", "Test prompt", api_base="FOO", api_key_var="TEST_API_KEY"
        )
        async_openai_mock.assert_called_once_with(
            api_key="test-key-12345",
            base_url="FOO",
        )
        openai_chat_completions_model.assert_called_once()


def test_run_openai_environment_error():
    with patch.dict(os.environ, {}, clear=True):
        with pytest.raises(KeyError, match="MISSING_KEY"):
            run_openai_agent(
                "test-model", "Test prompt", api_base="FOO", api_key_var="MISSING_KEY"
            )


def test_run_openai_multiagent():
    mock_agent = MagicMock()
    mock_function_tool = MagicMock()

    with (
        patch("surf_spot_finder.agents.openai.Agent", mock_agent),
        patch("surf_spot_finder.agents.openai.Runner", MagicMock()),
        patch("surf_spot_finder.agents.openai.function_tool", mock_function_tool),
    ):
        run_openai_multi_agent("gpt-4o", "Test prompt")
        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Interact with the user by showing information and asking for verification.",
            name="user-verification-agent",
            tools=[
                mock_function_tool(show_plan),
                mock_function_tool(ask_user_verification),
            ],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Find relevant information about the provided task by using your tools.",
            name="search-web-agent",
            tools=[mock_function_tool(search_web), mock_function_tool(visit_webpage)],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions="Communicate the final answer to the user.",
            name="communication-agent",
            tools=[mock_function_tool(show_final_answer)],
        )

        mock_agent.assert_any_call(
            model="gpt-4o",
            instructions=MULTI_AGENT_SYSTEM_PROMPT,
            name="surf-spot-finder",
            # TODO: add more elaborated checks
            handoffs=ANY,
            tools=ANY,
        )