Merge pull request #6 from andrewyng/start-testing

Start testing

Author: standsleeping

.github/workflows/run_pytest.yml

@@ -0,0 +1,24 @@
+name: Lint
+
+on: [push, pull_request]
+
+jobs:
+  build_and_test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [ "3.10", "3.11", "3.12" ]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install poetry
+          poetry install
+      - name: Test with pytest
+        run: poetry run pytest
+

tests/test_anthropic_interface.py

@@ -0,0 +1,46 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from aimodels.providers.anthropic_interface import AnthropicInterface
+
+
+@pytest.fixture(autouse=True)
+def set_api_key_env_var(monkeypatch):
+    """Fixture to set environment variables for tests."""
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "test-api-key")
+
+
+def test_anthropic_interface():
+    """High-level test that the interface is initialized and chat completions are requested successfully."""
+
+    user_greeting = "Hello!"
+    message_history = [{"role": "user", "content": user_greeting}]
+    selected_model = "our-favorite-model"
+    chosen_temperature = 0.75
+    response_text_content = "mocked-text-response-from-model"
+
+    interface = AnthropicInterface()
+    mock_response = MagicMock()
+    mock_response.content = [MagicMock()]
+    mock_response.content[0].text = response_text_content
+
+    with patch.object(
+        interface.anthropic_client.messages, "create", return_value=mock_response
+    ) as mock_create:
+        response = interface.chat_completion_create(
+            messages=message_history,
+            model=selected_model,
+            temperature=chosen_temperature,
+        )
+
+        transformed_message_history = [
+            {"role": "user", "content": [{"type": "text", "text": user_greeting}]},
+        ]
+
+        mock_create.assert_called_with(
+            messages=transformed_message_history,
+            model=selected_model,
+            temperature=chosen_temperature,
+            max_tokens=4096,
+        )
+
+        assert response.choices[0].message.content == response_text_content

tests/test_multi_fm_client.py

@@ -0,0 +1,45 @@
+import pytest
+from aimodels.client.multi_fm_client import MultiFMClient, AnthropicInterface
+
+
+def test_get_provider_interface_with_new_instance():
+    """Test that get_provider_interface creates a new instance of the interface."""
+    client = MultiFMClient()
+    interface, model_name = client.get_provider_interface("anthropic:some-model:v1")
+    assert isinstance(interface, AnthropicInterface)
+    assert model_name == "some-model:v1"
+    assert client.all_interfaces["anthropic"] == interface
+
+
+def test_get_provider_interface_with_existing_instance():
+    """Test that get_provider_interface returns an existing instance of the interface, if already created."""
+    client = MultiFMClient()
+
+    # New interface instance
+    new_instance, _ = client.get_provider_interface("anthropic:some-model:v2")
+
+    # Call twice, get same instance back
+    same_instance, _ = client.get_provider_interface("anthropic:some-model:v2")
+
+    assert new_instance is same_instance
+
+
+def test_get_provider_interface_with_invalid_format():
+    client = MultiFMClient()
+
+    with pytest.raises(ValueError) as exc_info:
+        client.get_provider_interface("invalid-model-no-colon")
+
+    assert "Expected ':' in model identifier" in str(exc_info.value)
+
+
+def test_get_provider_interface_with_unknown_interface():
+    client = MultiFMClient()
+
+    with pytest.raises(Exception) as exc_info:
+        client.get_provider_interface("unknown-interface:some-model")
+
+    assert (
+        "Could not find factory to create interface for provider 'unknown-interface'"
+        in str(exc_info.value)
+    )