Fixed test.

aana/tests/db/datastore/test_task_repo.py

@@ -45,44 +45,49 @@ def test_get_unprocessed_tasks(db_session):
     """Test fetching unprocessed tasks."""
     task_repo = TaskRepository(db_session)
 
-    # Remove all existing tasks
-    db_session.query(TaskEntity).delete()
-    db_session.commit()
-
-    # Create sample tasks with different statuses
-    now = datetime.now()  # noqa: DTZ005
-
-    task1 = TaskEntity(
-        endpoint="/test1",
-        data={"test": "data1"},
-        status=TaskStatus.CREATED,
-        priority=1,
-        created_at=now - timedelta(hours=10),
-    )
-    task2 = TaskEntity(
-        endpoint="/test2",
-        data={"test": "data2"},
-        status=TaskStatus.NOT_FINISHED,
-        priority=2,
-        created_at=now - timedelta(hours=1),
-    )
-    task3 = TaskEntity(
-        endpoint="/test3",
-        data={"test": "data3"},
-        status=TaskStatus.COMPLETED,
-        priority=3,
-        created_at=now - timedelta(hours=2),
-    )
-    task4 = TaskEntity(
-        endpoint="/test4",
-        data={"test": "data4"},
-        status=TaskStatus.CREATED,
-        priority=2,
-        created_at=now - timedelta(hours=3),
-    )
-
-    db_session.add_all([task1, task2, task3, task4])
-    db_session.commit()
+    def _create_sample_tasks():
+        # Remove all existing tasks
+        db_session.query(TaskEntity).delete()
+        db_session.commit()
+
+        # Create sample tasks with different statuses
+        now = datetime.now()  # noqa: DTZ005
+
+        task1 = TaskEntity(
+            endpoint="/test1",
+            data={"test": "data1"},
+            status=TaskStatus.CREATED,
+            priority=1,
+            created_at=now - timedelta(hours=10),
+        )
+        task2 = TaskEntity(
+            endpoint="/test2",
+            data={"test": "data2"},
+            status=TaskStatus.NOT_FINISHED,
+            priority=2,
+            created_at=now - timedelta(hours=1),
+        )
+        task3 = TaskEntity(
+            endpoint="/test3",
+            data={"test": "data3"},
+            status=TaskStatus.COMPLETED,
+            priority=3,
+            created_at=now - timedelta(hours=2),
+        )
+        task4 = TaskEntity(
+            endpoint="/test4",
+            data={"test": "data4"},
+            status=TaskStatus.CREATED,
+            priority=2,
+            created_at=now - timedelta(hours=3),
+        )
+
+        db_session.add_all([task1, task2, task3, task4])
+        db_session.commit()
+        return task1, task2, task3, task4
+
+    # Create sample tasks
+    task1, task2, task3, task4 = _create_sample_tasks()
 
     # Fetch unprocessed tasks without any limit
     unprocessed_tasks = task_repo.fetch_unprocessed_tasks()
@@ -98,6 +103,9 @@ def test_get_unprocessed_tasks(db_session):
     assert unprocessed_tasks[1].id == task2.id  # Same priority, but a newer task
     assert unprocessed_tasks[2].id == task1.id  # Lowest priority
 
+    # Create sample tasks
+    task1, task2, task3, task4 = _create_sample_tasks()
+
     # Fetch unprocessed tasks with a limit
     limited_tasks = task_repo.fetch_unprocessed_tasks(limit=2)