fix failing tests

This commit is contained in:
Joshua Vial 2023-12-11 22:37:23 +13:00
parent 9ceaf97f08
commit 90d5071709
2 changed files with 16 additions and 13 deletions

View file

@ -29,19 +29,19 @@ class TestModels(unittest.TestCase):
def test_openrouter_model_properties(self):
client = MagicMock()
client.models.list.return_value = {
"data": [
{
"id": "openai/gpt-4",
"object": "model",
"context_length": "8192",
"pricing": {"prompt": "0.00006", "completion": "0.00012"},
}
]
}
client.models.list.return_value = type(
"", (), {"data": client.models.list.return_value["data"]}
)()
class ModelData:
def __init__(self, id, object, context_length, pricing):
self.id = id
self.object = object
self.context_length = context_length
self.pricing = pricing
model_data = ModelData("openai/gpt-4", "model", "8192", {"prompt": "0.00006", "completion": "0.00012"})
class ModelList:
def __init__(self, data):
self.data = data
client.models.list.return_value = ModelList([model_data])
model = OpenRouterModel(client, "gpt-4")
self.assertEqual(model.name, "openai/gpt-4")