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