mock gemma

allegro · Mar 6, 2024 · b987dba · b987dba
1 parent 70efeed
commit b987dba
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -37,7 +37,8 @@ def models():
     event_loop = asyncio.new_event_loop()
 
     with (patch("allms.models.vertexai_palm.CustomVertexAI", VertexAIMock),
-          patch("allms.models.vertexai_gemini.CustomVertexAI", VertexAIMock)):
+          patch("allms.models.vertexai_gemini.CustomVertexAI", VertexAIMock),
+          patch("allms.models.vertexai_gemma.VertexAIModelGardenWrapper", VertexAIMock)):
         return {
                 "azure_open_ai": AzureOpenAIModel(
                     config=AzureOpenAIConfiguration(