From e1a0b3ffa2aa6e5a7c068cdc559ca126f3b57b4c Mon Sep 17 00:00:00 2001 From: Heiner Lohaus Date: Wed, 14 Feb 2024 09:21:57 +0100 Subject: Add unitests for the client Fix: async generator ignored GeneratorExit Fix: ResourceWarning: unclosed event loop --- etc/unittest/asyncio.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) (limited to 'etc/unittest/asyncio.py') diff --git a/etc/unittest/asyncio.py b/etc/unittest/asyncio.py index e886c43a..57a1fb7d 100644 --- a/etc/unittest/asyncio.py +++ b/etc/unittest/asyncio.py @@ -8,6 +8,7 @@ import unittest import g4f from g4f import ChatCompletion +from g4f.client import Client from .mocks import ProviderMock, AsyncProviderMock, AsyncGeneratorProviderMock DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] @@ -24,11 +25,16 @@ class TestChatCompletion(unittest.TestCase): def test_create(self): result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock) - self.assertEqual("Mock",result) + self.assertEqual("Mock", result) def test_create_generator(self): result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock) - self.assertEqual("Mock",result) + self.assertEqual("Mock", result) + + def test_await_callback(self): + client = Client(provider=AsyncGeneratorProviderMock) + response = client.chat.completions.create(DEFAULT_MESSAGES, "", max_tokens=0) + self.assertEqual("Mock", response.choices[0].message.content) class TestChatCompletionAsync(unittest.IsolatedAsyncioTestCase): -- cgit v1.2.3