From 293c565acfa8af6b230d65ee4c088c940a7cf400 Mon Sep 17 00:00:00 2001
From: "jenn.muengtaweepongsa" <jenn.muengtaweepongsa@sentry.io>
Date: Fri, 6 Sep 2024 14:02:53 -0700
Subject: [PATCH] fix all tests

---
 tests/automation/autofix/test_autofix_context.py |  5 +++--
 tests/automation/autofix/test_autofix_state.py   | 15 ++++-----------
 tests/test_celery.py                             |  3 ++-
 3 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/tests/automation/autofix/test_autofix_context.py b/tests/automation/autofix/test_autofix_context.py
index 64999863..2de88406 100644
--- a/tests/automation/autofix/test_autofix_context.py
+++ b/tests/automation/autofix/test_autofix_context.py
@@ -25,7 +25,7 @@
     StacktraceFrame,
     ThreadDetails,
 )
-from seer.automation.state import LocalMemoryState
+from seer.automation.state import DbStateRunTypes, LocalMemoryState
 from seer.automation.summarize.issue import IssueSummary
 from seer.db import DbIssueSummary, DbPrIdToAutofixRunIdMapping, Session
 
@@ -170,7 +170,8 @@ def setUp(self):
                     repos=[],
                     issue=IssueDetails(id=0, title="", events=[error_event], short_id="ISSUE_1"),
                 ),
-            )
+            ),
+            type=DbStateRunTypes.AUTOFIX,
         )
         self.autofix_context = AutofixContext(self.state, MagicMock(), MagicMock())
         self.autofix_context.get_org_slug = MagicMock(return_value="slug")
diff --git a/tests/automation/autofix/test_autofix_state.py b/tests/automation/autofix/test_autofix_state.py
index 285d7949..2601f223 100644
--- a/tests/automation/autofix/test_autofix_state.py
+++ b/tests/automation/autofix/test_autofix_state.py
@@ -4,6 +4,7 @@
 
 from seer.automation.autofix.models import AutofixContinuation
 from seer.automation.autofix.state import ContinuationState
+from seer.automation.state import DbStateRunTypes
 from seer.db import DbRunState
 
 
@@ -13,19 +14,11 @@ def mock_session(self):
         with patch("seer.automation.autofix.state.Session") as mock:
             yield mock
 
-    def test_from_id(self):
-        with patch("seer.automation.state.DbState.from_id") as mock_from_id:
-            mock_from_id.return_value = MagicMock(spec=ContinuationState)
-            result = ContinuationState.from_id(1, AutofixContinuation)
-
-            mock_from_id.assert_called_once_with(1, AutofixContinuation)
-            assert isinstance(result, ContinuationState)
-
     def test_set(self, mock_session):
         mock_session_instance = MagicMock()
         mock_session.return_value.__enter__.return_value = mock_session_instance
 
-        state = ContinuationState(id=1, model=AutofixContinuation)
+        state = ContinuationState(id=1, model=AutofixContinuation, type=DbStateRunTypes.AUTOFIX)
         mock_autofix_continuation = MagicMock(spec=AutofixContinuation)
         mock_autofix_continuation.model_dump.return_value = {"key": "value"}
         mock_autofix_continuation.updated_at = "2023-01-01"
@@ -48,7 +41,7 @@ def test_set(self, mock_session):
     def test_get(self):
         with patch("seer.automation.state.DbState.get") as mock_get:
             mock_get.return_value = MagicMock(spec=AutofixContinuation)
-            state = ContinuationState(id=1, model=AutofixContinuation)
+            state = ContinuationState(id=1, model=AutofixContinuation, type=DbStateRunTypes.AUTOFIX)
             result = state.get()
 
             mock_get.assert_called_once()
@@ -59,7 +52,7 @@ def test_update(self):
             mock_context = MagicMock()
             mock_update.return_value.__enter__.return_value = mock_context
 
-            state = ContinuationState(id=1, model=AutofixContinuation)
+            state = ContinuationState(id=1, model=AutofixContinuation, type=DbStateRunTypes.AUTOFIX)
             with state.update() as context:
                 assert context == mock_context
 
diff --git a/tests/test_celery.py b/tests/test_celery.py
index 73f7e4e7..b10570ed 100644
--- a/tests/test_celery.py
+++ b/tests/test_celery.py
@@ -1,5 +1,6 @@
 from celery import Celery
 
+import celery_app.tasks
 from celery_app.app import celery_app, setup_celery_entrypoint
 from celery_app.config import CeleryConfig
 from seer.dependency_injection import resolve
@@ -15,7 +16,7 @@ def test_detected_celery_jobs():
             "seer.automation.autofix.steps.steps.autofix_parallelized_conditional_step_task",
             "seer.automation.autofix.tasks.run_autofix_evaluation_on_item",
             "seer.automation.autofix.tasks.check_and_mark_recent_autofix_runs",
-            "seer.automation.autofix.tasks.delete_data_for_ttl",
+            "seer.automation.tasks.delete_data_for_ttl",
             "seer.smoke_test.smoke_test",
         ]
     )