44 lines
1.8 KiB
Python
44 lines
1.8 KiB
Python
|
from unittest.mock import MagicMock, patch
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from rein_openweights.openrouter_client import OpenRouterClient, OpenRouterError
|
||
|
|
|
||
|
|
|
||
|
|
def test_chat_sends_bearer_auth_and_model():
|
||
|
|
client = OpenRouterClient(api_key="sk-test", model="meta-llama/llama-3.1-70b-instruct")
|
||
|
|
fake_response = MagicMock(status_code=200)
|
||
|
|
fake_response.json.return_value = {"choices": [{"message": {"content": "hi"}}]}
|
||
|
|
|
||
|
|
with patch("rein_openweights.openrouter_client.httpx.post", return_value=fake_response) as post:
|
||
|
|
data = client.chat([{"role": "user", "content": "hello"}])
|
||
|
|
|
||
|
|
assert data["choices"][0]["message"]["content"] == "hi"
|
||
|
|
_, kwargs = post.call_args
|
||
|
|
assert kwargs["headers"]["Authorization"] == "Bearer sk-test"
|
||
|
|
assert kwargs["json"]["model"] == "meta-llama/llama-3.1-70b-instruct"
|
||
|
|
assert "tools" not in kwargs["json"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_chat_includes_tools_when_provided():
|
||
|
|
client = OpenRouterClient(api_key="sk-test", model="m")
|
||
|
|
fake_response = MagicMock(status_code=200)
|
||
|
|
fake_response.json.return_value = {"choices": [{"message": {}}]}
|
||
|
|
tools = [{"type": "function", "function": {"name": "read_file"}}]
|
||
|
|
|
||
|
|
with patch("rein_openweights.openrouter_client.httpx.post", return_value=fake_response) as post:
|
||
|
|
client.chat([{"role": "user", "content": "x"}], tools=tools)
|
||
|
|
|
||
|
|
_, kwargs = post.call_args
|
||
|
|
assert kwargs["json"]["tools"] == tools
|
||
|
|
assert kwargs["json"]["tool_choice"] == "auto"
|
||
|
|
|
||
|
|
|
||
|
|
def test_chat_raises_on_http_error():
|
||
|
|
client = OpenRouterClient(api_key="sk-test", model="m")
|
||
|
|
fake_response = MagicMock(status_code=401, text="unauthorized")
|
||
|
|
|
||
|
|
with patch("rein_openweights.openrouter_client.httpx.post", return_value=fake_response):
|
||
|
|
with pytest.raises(OpenRouterError, match="401"):
|
||
|
|
client.chat([{"role": "user", "content": "x"}])
|