🐛✅ Fix fixture for tgis tests

Signed-off-by: gkumbhat <[email protected]>
opendatahub-io · May 21, 2024 · 3fea23c · 3fea23c
1 parent b013b50
commit 3fea23c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py
@@ -208,7 +208,7 @@ def Tokenize(self, request, **kwargs):
         return StubTGISClient.tokenize(request)
 
     @staticmethod
-    def unary_generate(request):
+    def unary_generate(request, **kwargs):
         fake_response = mock.Mock()
         fake_result = mock.Mock()
         fake_result.stop_reason = 5
@@ -229,7 +229,7 @@ def unary_generate(request):
         return fake_response
 
     @staticmethod
-    def stream_generate(request):
+    def stream_generate(request, **kwargs):
         fake_stream = mock.Mock()
         fake_stream.stop_reason = 5
         fake_stream.generated_token_count = 1
@@ -250,7 +250,7 @@ def stream_generate(request):
             yield fake_stream
 
     @staticmethod
-    def tokenize(request):
+    def tokenize(request, **kwargs):
         fake_response = mock.Mock()
         fake_result = mock.Mock()
         fake_result.token_count = 1

diff --git a/tests/toolkit/text_generation/test_tgis_utils.py b/tests/toolkit/text_generation/test_tgis_utils.py
@@ -54,22 +54,19 @@ def _maybe_raise(self, error_type: Type[grpc.RpcError], *args):
             )
 
     def Generate(
-        self,
-        request: generation_pb2.BatchedGenerationRequest,
+        self, request: generation_pb2.BatchedGenerationRequest, **kwargs
     ) -> generation_pb2.BatchedGenerationResponse:
         self._maybe_raise(grpc._channel._InactiveRpcError)
         return generation_pb2.BatchedGenerationResponse()
 
     def GenerateStream(
-        self,
-        request: generation_pb2.SingleGenerationRequest,
+        self, request: generation_pb2.SingleGenerationRequest, **kwargs
     ) -> Iterable[generation_pb2.GenerationResponse]:
         self._maybe_raise(grpc._channel._MultiThreadedRendezvous, None, None, None)
         yield generation_pb2.GenerationResponse()
 
     def Tokenize(
-        self,
-        request: generation_pb2.BatchedTokenizeRequest,
+        self, request: generation_pb2.BatchedTokenizeRequest, **kwargs
     ) -> generation_pb2.BatchedTokenizeResponse:
         self._maybe_raise(grpc._channel._InactiveRpcError)
         return generation_pb2.BatchedTokenizeResponse()