44from fastapi .testclient import TestClient
55
66from khoj .utils import state
7- from tests .helpers import ChatModelFactory
7+ from tests .helpers import AiModelApiFactory , ChatModelFactory , get_chat_api_key
88
99
1010@pytest .fixture (autouse = True )
@@ -19,7 +19,9 @@ def setup_scheduler():
1919def create_test_automation (client : TestClient ) -> str :
2020 """Helper function to create a test automation and return its ID."""
2121 state .anonymous_mode = True
22- ChatModelFactory (name = "gpt-4o-mini" , model_type = "openai" )
22+ ChatModelFactory (
23+ name = "gemini-2.0-flash" , model_type = "google" , ai_model_api = AiModelApiFactory (api_key = get_chat_api_key ("google" ))
24+ )
2325 params = {
2426 "q" : "test automation" ,
2527 "crontime" : "0 0 * * *" ,
@@ -34,7 +36,9 @@ def test_create_automation(client: TestClient):
3436 """Test that creating an automation works as expected."""
3537 # Arrange
3638 state .anonymous_mode = True
37- ChatModelFactory (name = "gpt-4o-mini" , model_type = "openai" )
39+ ChatModelFactory (
40+ name = "gemini-2.0-flash" , model_type = "google" , ai_model_api = AiModelApiFactory (api_key = get_chat_api_key ("google" ))
41+ )
3842 params = {
3943 "q" : "test automation" ,
4044 "crontime" : "0 0 * * *" ,
@@ -51,6 +55,7 @@ def test_create_automation(client: TestClient):
5155
5256
5357@pytest .mark .django_db (transaction = True )
58+ @pytest .mark .skipif (get_chat_api_key ("google" ) is None , reason = "Requires GEMINI_API_KEY to be set" )
5459def test_get_automations (client : TestClient ):
5560 """Test that getting a list of automations works."""
5661 automation_id = create_test_automation (client )
@@ -67,6 +72,7 @@ def test_get_automations(client: TestClient):
6772
6873
6974@pytest .mark .django_db (transaction = True )
75+ @pytest .mark .skipif (get_chat_api_key ("google" ) is None , reason = "Requires GEMINI_API_KEY to be set" )
7076def test_delete_automation (client : TestClient ):
7177 """Test that deleting an automation works."""
7278 automation_id = create_test_automation (client )
@@ -85,6 +91,7 @@ def test_delete_automation(client: TestClient):
8591
8692
8793@pytest .mark .django_db (transaction = True )
94+ @pytest .mark .skipif (get_chat_api_key ("google" ) is None , reason = "Requires GEMINI_API_KEY to be set" )
8895def test_edit_automation (client : TestClient ):
8996 """Test that editing an automation works."""
9097 automation_id = create_test_automation (client )
@@ -111,6 +118,7 @@ def test_edit_automation(client: TestClient):
111118
112119
113120@pytest .mark .django_db (transaction = True )
121+ @pytest .mark .skipif (get_chat_api_key ("google" ) is None , reason = "Requires GEMINI_API_KEY to be set" )
114122def test_trigger_automation (client : TestClient ):
115123 """Test that triggering an automation works."""
116124 automation_id = create_test_automation (client )
0 commit comments