Skip to content

test_openai_server

ipw.tests.agents.mcp.test_openai_server

Tests for agents/mcp/openai_server.py -- OpenAIMCPServer pricing.

TestOpenAIMCPServerPricing

Test OpenAI server pricing calculations with mocked openai.

Source code in intelligence-per-watt/src/ipw/tests/agents/mcp/test_openai_server.py
class TestOpenAIMCPServerPricing:
    """Test OpenAI server pricing calculations with mocked openai."""

    def test_pricing_calculation_gpt4o(self) -> None:
        server = _make_server("gpt-4o")

        assert server.pricing is not None
        assert server.pricing["input"] == 2.50
        assert server.pricing["output"] == 10.00

        cost = calculate_cost("openai", "gpt-4o", 1000, 500)
        expected = (1000 / 1_000_000) * 2.50 + (500 / 1_000_000) * 10.00
        assert cost == pytest.approx(expected)

    def test_unknown_model_falls_back_to_gpt4o_pricing(self) -> None:
        server = _make_server("totally-unknown-model")
        assert server.pricing["input"] == 2.50
        assert server.pricing["output"] == 10.00

    def test_name_includes_model(self) -> None:
        server = _make_server("gpt-4o-mini")
        assert server.name == "openai:gpt-4o-mini"

    def test_execute_impl_with_mocked_stream(self) -> None:
        _MOCK_OPENAI.OpenAI.reset_mock()
        _MOCK_OPENAI.OpenAIError = type("OpenAIError", (Exception,), {})

        mock_client = MagicMock()
        _MOCK_OPENAI.OpenAI.return_value = mock_client

        # Build mock streaming response
        mock_chunk1 = MagicMock()
        mock_chunk1.choices = [MagicMock()]
        mock_chunk1.choices[0].delta.content = "Hello"
        mock_chunk1.usage = None

        mock_chunk2 = MagicMock()
        mock_chunk2.choices = [MagicMock()]
        mock_chunk2.choices[0].delta.content = " world"
        mock_chunk2.usage = None

        mock_chunk3 = MagicMock()
        mock_chunk3.choices = []
        mock_usage = MagicMock()
        mock_usage.prompt_tokens = 10
        mock_usage.completion_tokens = 5
        mock_usage.total_tokens = 15
        mock_chunk3.usage = mock_usage

        mock_client.chat.completions.create.return_value = iter(
            [mock_chunk1, mock_chunk2, mock_chunk3]
        )

        server = OpenAIMCPServer(model_name="gpt-4o", api_key="test-key")
        result = server._execute_impl("test prompt")

        assert result.content == "Hello world"
        assert result.usage["prompt_tokens"] == 10
        assert result.usage["completion_tokens"] == 5
        assert result.cost_usd > 0