44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import pytest
|
|
import sys
|
|
import os
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
# Ensure project root is in path
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
import ai_client
|
|
|
|
def test_agent_capabilities_listing() -> None:
|
|
# Mock credentials
|
|
with patch("ai_client._load_credentials") as mock_creds:
|
|
mock_creds.return_value = {"gemini": {"api_key": "fake-key"}}
|
|
|
|
# Mock the google-genai Client and models.list
|
|
with patch("google.genai.Client") as mock_client_class:
|
|
mock_client = mock_client_class.return_value
|
|
|
|
# Create a list of mock models
|
|
mock_models = []
|
|
model_names = [
|
|
"models/gemini-2.0-flash",
|
|
"models/gemini-2.0-flash-lite",
|
|
"models/gemini-1.5-pro",
|
|
"models/gemini-1.5-flash",
|
|
"models/gemini-pro",
|
|
"models/gemini-ultra"
|
|
]
|
|
for name in model_names:
|
|
m = MagicMock()
|
|
m.name = name
|
|
mock_models.append(m)
|
|
|
|
mock_client.models.list.return_value = mock_models
|
|
|
|
models = ai_client.list_models("gemini")
|
|
|
|
assert len(models) >= 6
|
|
assert "gemini-2.0-flash" in models
|
|
# Check that it stripped "models/" prefix
|
|
for m in models:
|
|
assert not m.startswith("models/")
|