1  import pytest
  2  from apscheduler.schedulers.background import BackgroundScheduler
  3  from django_apscheduler.jobstores import DjangoJobStore
  4  from fastapi.testclient import TestClient
  5  
  6  from khoj.utils import state
  7  from tests.helpers import AiModelApiFactory, ChatModelFactory, get_chat_api_key
  8  
  9  
 10  @pytest.fixture(autouse=True)
 11  def setup_scheduler():
 12      state.scheduler = BackgroundScheduler()
 13      state.scheduler.add_jobstore(DjangoJobStore(), "default")
 14      state.scheduler.start()
 15      yield
 16      state.scheduler.shutdown()
 17  
 18  
 19  def create_test_automation(client: TestClient) -> str:
 20      """Helper function to create a test automation and return its ID."""
 21      state.anonymous_mode = True
 22      ChatModelFactory(
 23          name="gemini-2.5-flash", model_type="google", ai_model_api=AiModelApiFactory(api_key=get_chat_api_key("google"))
 24      )
 25      params = {
 26          "q": "test automation",
 27          "crontime": "0 0 * * *",
 28      }
 29      response = client.post("/api/automation", params=params)
 30      assert response.status_code == 200
 31      return response.json()["id"]
 32  
 33  
 34  @pytest.mark.django_db(transaction=True)
 35  @pytest.mark.skipif(not get_chat_api_key("google"), reason="Requires GEMINI_API_KEY to be set")
 36  def test_create_automation(client: TestClient):
 37      """Test that creating an automation works as expected."""
 38      # Arrange
 39      state.anonymous_mode = True
 40      ChatModelFactory(
 41          name="gemini-2.5-flash", model_type="google", ai_model_api=AiModelApiFactory(api_key=get_chat_api_key("google"))
 42      )
 43      params = {
 44          "q": "test automation",
 45          "crontime": "0 0 * * *",
 46      }
 47  
 48      # Act
 49      response = client.post("/api/automation", params=params)
 50  
 51      # Assert
 52      assert response.status_code == 200
 53      response_json = response.json()
 54      assert response_json["scheduling_request"] == "test automation"
 55      assert response_json["crontime"] == "0 0 * * *"
 56  
 57  
 58  @pytest.mark.django_db(transaction=True)
 59  @pytest.mark.skipif(not get_chat_api_key("google"), reason="Requires GEMINI_API_KEY to be set")
 60  def test_get_automations(client: TestClient):
 61      """Test that getting a list of automations works."""
 62      automation_id = create_test_automation(client)
 63  
 64      # Act
 65      response = client.get("/api/automation")
 66  
 67      # Assert
 68      assert response.status_code == 200
 69      automations = response.json()
 70      assert isinstance(automations, list)
 71      assert len(automations) > 0
 72      assert any(a["id"] == automation_id for a in automations)
 73  
 74  
 75  @pytest.mark.django_db(transaction=True)
 76  @pytest.mark.skipif(not get_chat_api_key("google"), reason="Requires GEMINI_API_KEY to be set")
 77  def test_delete_automation(client: TestClient):
 78      """Test that deleting an automation works."""
 79      automation_id = create_test_automation(client)
 80  
 81      # Act
 82      response = client.delete(f"/api/automation?automation_id={automation_id}")
 83  
 84      # Assert
 85      assert response.status_code == 200
 86  
 87      # Verify it's gone
 88      response = client.get("/api/automation")
 89      assert response.status_code == 200
 90      automations = response.json()
 91      assert not any(a["id"] == automation_id for a in automations)
 92  
 93  
 94  @pytest.mark.django_db(transaction=True)
 95  @pytest.mark.skipif(not get_chat_api_key("google"), reason="Requires GEMINI_API_KEY to be set")
 96  def test_edit_automation(client: TestClient):
 97      """Test that editing an automation works."""
 98      automation_id = create_test_automation(client)
 99  
100      edit_params = {
101          "automation_id": automation_id,
102          "q": "edited automation",
103          "crontime": "0 1 * * *",
104          "subject": "edited subject",
105          "timezone": "UTC",
106      }
107  
108      # Act
109      response = client.put("/api/automation", params=edit_params)
110  
111      # Assert
112      if response.status_code != 200:
113          print(response.text)
114      assert response.status_code == 200
115      edited_automation = response.json()
116      assert edited_automation["scheduling_request"] == "edited automation"
117      assert edited_automation["crontime"] == "0 1 * * *"
118      assert edited_automation["subject"] == "edited subject"
119  
120  
121  @pytest.mark.django_db(transaction=True)
122  @pytest.mark.skipif(not get_chat_api_key("google"), reason="Requires GEMINI_API_KEY to be set")
123  def test_trigger_automation(client: TestClient):
124      """Test that triggering an automation works."""
125      automation_id = create_test_automation(client)
126  
127      # Act
128      response = client.post(f"/api/automation/trigger?automation_id={automation_id}")
129  
130      # Assert
131      assert response.status_code == 200
132      # NOTE: We are not testing the execution of the triggered job itself,
133      # as that would require a more complex test setup with mocking.
134      # A 200 response is sufficient to indicate that the trigger was received.
135      assert response.text == "Automation triggered"