oops

Author: khvn26

src/common/test_tools/__init__.py

@@ -1,6 +1,11 @@
-from common.test_tools.types import AssertMetricFixture, SnapshotFixture
+from common.test_tools.types import (
+    AssertMetricFixture,
+    RunTasksFixture,
+    SnapshotFixture,
+)
 
 __all__ = (
     "AssertMetricFixture",
+    "RunTasksFixture",
     "SnapshotFixture",
 )

src/common/test_tools/types.py

@@ -1,8 +1,11 @@
 from pathlib import Path
-from typing import Protocol
+from typing import TYPE_CHECKING, Protocol
 
 import pytest
 
+if TYPE_CHECKING:
+    from task_processor.models import TaskRun
+
 
 class AssertMetricFixture(Protocol):
     def __call__(
@@ -14,6 +17,13 @@ def __call__(
     ) -> None: ...
 
 
+class RunTasksFixture(Protocol):
+    def __call__(
+        self,
+        num_tasks: int,
+    ) -> "list[TaskRun]": ...
+
+
 class SnapshotFixture(Protocol):
     def __call__(self, name: str = "") -> "Snapshot": ...