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/mocks.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) (limited to 'etc/unittest/mocks.py') diff --git a/etc/unittest/mocks.py b/etc/unittest/mocks.py index 885bdaee..8a67aaf7 100644 --- a/etc/unittest/mocks.py +++ b/etc/unittest/mocks.py @@ -7,10 +7,10 @@ class ProviderMock(AbstractProvider): model, messages, stream, **kwargs ): yield "Mock" - + class AsyncProviderMock(AsyncProvider): working = True - + async def create_async( model, messages, **kwargs ): @@ -18,16 +18,25 @@ class AsyncProviderMock(AsyncProvider): class AsyncGeneratorProviderMock(AsyncGeneratorProvider): working = True - + async def create_async_generator( model, messages, stream, **kwargs ): yield "Mock" - + class ModelProviderMock(AbstractProvider): working = True def create_completion( model, messages, stream, **kwargs ): - yield model \ No newline at end of file + yield model + +class YieldProviderMock(AsyncGeneratorProvider): + working = True + + async def create_async_generator( + model, messages, stream, **kwargs + ): + for message in messages: + yield message["content"] \ No newline at end of file -- cgit v1.2.3