Fix to run new automation api tests in ci

Author: debanjum

.github/workflows/test.yml

@@ -92,5 +92,6 @@ jobs:
           POSTGRES_USER: postgres
           POSTGRES_PASSWORD: postgres
           POSTGRES_DB: postgres
+          GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
         run: pytest
         timeout-minutes: 10

tests/test_api_automation.py

@@ -4,7 +4,7 @@
 from fastapi.testclient import TestClient
 
 from khoj.utils import state
-from tests.helpers import ChatModelFactory
+from tests.helpers import AiModelApiFactory, ChatModelFactory, get_chat_api_key
 
 
 @pytest.fixture(autouse=True)
@@ -19,7 +19,9 @@ def setup_scheduler():
 def create_test_automation(client: TestClient) -> str:
     """Helper function to create a test automation and return its ID."""
     state.anonymous_mode = True
-    ChatModelFactory(name="gpt-4o-mini", model_type="openai")
+    ChatModelFactory(
+        name="gemini-2.0-flash", model_type="google", ai_model_api=AiModelApiFactory(api_key=get_chat_api_key("google"))
+    )
     params = {
         "q": "test automation",
         "crontime": "0 0 * * *",
@@ -34,7 +36,9 @@ def test_create_automation(client: TestClient):
     """Test that creating an automation works as expected."""
     # Arrange
     state.anonymous_mode = True
-    ChatModelFactory(name="gpt-4o-mini", model_type="openai")
+    ChatModelFactory(
+        name="gemini-2.0-flash", model_type="google", ai_model_api=AiModelApiFactory(api_key=get_chat_api_key("google"))
+    )
     params = {
         "q": "test automation",
         "crontime": "0 0 * * *",
@@ -51,6 +55,7 @@ def test_create_automation(client: TestClient):
 
 
 @pytest.mark.django_db(transaction=True)
+@pytest.mark.skipif(get_chat_api_key("google") is None, reason="Requires GEMINI_API_KEY to be set")
 def test_get_automations(client: TestClient):
     """Test that getting a list of automations works."""
     automation_id = create_test_automation(client)
@@ -67,6 +72,7 @@ def test_get_automations(client: TestClient):
 
 
 @pytest.mark.django_db(transaction=True)
+@pytest.mark.skipif(get_chat_api_key("google") is None, reason="Requires GEMINI_API_KEY to be set")
 def test_delete_automation(client: TestClient):
     """Test that deleting an automation works."""
     automation_id = create_test_automation(client)
@@ -85,6 +91,7 @@ def test_delete_automation(client: TestClient):
 
 
 @pytest.mark.django_db(transaction=True)
+@pytest.mark.skipif(get_chat_api_key("google") is None, reason="Requires GEMINI_API_KEY to be set")
 def test_edit_automation(client: TestClient):
     """Test that editing an automation works."""
     automation_id = create_test_automation(client)
@@ -111,6 +118,7 @@ def test_edit_automation(client: TestClient):
 
 
 @pytest.mark.django_db(transaction=True)
+@pytest.mark.skipif(get_chat_api_key("google") is None, reason="Requires GEMINI_API_KEY to be set")
 def test_trigger_automation(client: TestClient):
     """Test that triggering an automation works."""
     automation_id = create_test_automation(client)

tests/test_online_chat_actors.py

@@ -1,3 +1,4 @@
+import os
 from datetime import datetime
 
 import freezegun
@@ -20,7 +21,7 @@
 
 # Initialize variables for tests
 api_key = get_chat_api_key()
-if api_key is None:
+if api_key is None or not os.getenv("KHOJ_TEST_CHAT_PROVIDER"):
     pytest.skip(
         reason="Set OPENAI_API_KEY, GEMINI_API_KEY or ANTHROPIC_API_KEY environment variable to run tests below.",
         allow_module_level=True,

tests/test_online_chat_director.py

@@ -10,9 +10,9 @@
 
 # Initialize variables for tests
 api_key = get_chat_api_key()
-if api_key is None:
+if api_key is None or not os.getenv("KHOJ_TEST_CHAT_PROVIDER"):
     pytest.skip(
-        reason="Set OPENAI_API_KEY environment variable to run tests below. Get OpenAI API key from https://platform.openai.com/account/api-keys",
+        reason="Set OPENAI_API_KEY, GEMINI_API_KEY or ANTHROPIC_API_KEY environment variable to run tests below.",
         allow_module_level=True,
     )