summaryrefslogtreecommitdiffstats
path: root/etc/unittest/mocks.py
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--etc/unittest/mocks.py19
1 files changed, 14 insertions, 5 deletions
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