File size: 4,658 Bytes
39e315a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """
test_atomic_chat_provider.py
Run: pytest python/tests/test_atomic_chat_provider.py -v
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from atomic_chat_provider import (
atomic_chat,
list_atomic_chat_models,
check_atomic_chat_running,
)
@pytest.mark.asyncio
async def test_atomic_chat_running_true():
mock_response = MagicMock()
mock_response.status_code = 200
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.get = AsyncMock(return_value=mock_response)
result = await check_atomic_chat_running()
assert result is True
@pytest.mark.asyncio
async def test_atomic_chat_running_false_on_exception():
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.get = AsyncMock(side_effect=Exception("refused"))
result = await check_atomic_chat_running()
assert result is False
@pytest.mark.asyncio
async def test_list_models_returns_ids():
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": [{"id": "llama-3.1-8b"}, {"id": "mistral-7b"}],
}
mock_response.raise_for_status = MagicMock()
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.get = AsyncMock(return_value=mock_response)
models = await list_atomic_chat_models()
assert "llama-3.1-8b" in models
assert "mistral-7b" in models
@pytest.mark.asyncio
async def test_list_models_empty_on_failure():
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.get = AsyncMock(side_effect=Exception("down"))
models = await list_atomic_chat_models()
assert models == []
@pytest.mark.asyncio
async def test_atomic_chat_returns_anthropic_format():
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock()
mock_response.json.return_value = {
"id": "chatcmpl-abc123",
"choices": [{"message": {"content": "42 is the answer."}}],
"usage": {"prompt_tokens": 10, "completion_tokens": 8},
}
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.post = AsyncMock(return_value=mock_response)
result = await atomic_chat(
model="llama-3.1-8b",
messages=[{"role": "user", "content": "What is 6*7?"}],
)
assert result["type"] == "message"
assert result["role"] == "assistant"
assert "42" in result["content"][0]["text"]
assert result["usage"]["input_tokens"] == 10
assert result["usage"]["output_tokens"] == 8
@pytest.mark.asyncio
async def test_atomic_chat_prepends_system():
captured = {}
async def mock_post(url, json=None, **kwargs):
captured.update(json or {})
m = MagicMock()
m.raise_for_status = MagicMock()
m.json.return_value = {
"id": "chatcmpl-xyz",
"choices": [{"message": {"content": "ok"}}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
}
return m
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.post = mock_post
await atomic_chat(
model="llama-3.1-8b",
messages=[{"role": "user", "content": "Hi"}],
system="Be helpful.",
)
assert captured["messages"][0]["role"] == "system"
assert "helpful" in captured["messages"][0]["content"]
@pytest.mark.asyncio
async def test_atomic_chat_sends_correct_payload():
captured = {}
async def mock_post(url, json=None, **kwargs):
captured.update(json or {})
m = MagicMock()
m.raise_for_status = MagicMock()
m.json.return_value = {
"id": "chatcmpl-xyz",
"choices": [{"message": {"content": "ok"}}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
}
return m
with patch("atomic_chat_provider.httpx.AsyncClient") as MockClient:
MockClient.return_value.__aenter__.return_value.post = mock_post
await atomic_chat(
model="test-model",
messages=[{"role": "user", "content": "Test"}],
max_tokens=2048,
temperature=0.5,
)
assert captured["model"] == "test-model"
assert captured["max_tokens"] == 2048
assert captured["temperature"] == 0.5
assert captured["stream"] is False
|