1 files changed, 279 insertions, 0 deletions
diff --git a/tests/fixtures/models.py b/tests/fixtures/models.py
new file mode 100644
index 0000000..03bd8a2
--- /dev/null
+++ b/tests/fixtures/models.py
@@ -0,0 +1,279 @@
+from datetime import timezone, datetime
+
+import pytest
+
+from jb.models.event import MTurkEvent
+from generalresearchutils.pg_helper import PostgresConfig
+
+from datetime import datetime, timezone, timedelta
+from typing import Optional, TYPE_CHECKING, Callable, Generator
+from jb.managers.amt import AMTManager
+from jb.models.assignment import AssignmentStub, Assignment
+from generalresearchutils.currency import USDCent
+from jb.models.definitions import HitStatus, HitReviewStatus, AssignmentStatus
+from jb.models.hit import HitType, HitQuestion, Hit
+from tests import generate_amt_id
+
+if TYPE_CHECKING:
+    from jb.managers.hit import HitQuestionManager, HitTypeManager, HitManager
+    from jb.managers.assignment import AssignmentManager
+
+
+# --- MTurk Event ---
+
+
+@pytest.fixture
+def mturk_event(
+    amt_assignment_id: str, amt_hit_id: str, amt_hit_type_id: str
+) -> MTurkEvent:
+    now = datetime.now(tz=timezone.utc)
+    return MTurkEvent(
+        event_type="AssignmentSubmitted",
+        event_timestamp=now,
+        amt_assignment_id=amt_assignment_id,
+        amt_hit_type_id=amt_hit_type_id,
+        amt_hit_id=amt_hit_id,
+    )
+
+
+# --- Question ---
+
+
+@pytest.fixture
+def question() -> HitQuestion:
+    return HitQuestion(url="https://jamesbillings67.com/work/", height=1200)
+
+
+@pytest.fixture
+def question_record(hqm: "HitQuestionManager", question: HitQuestion) -> HitQuestion:
+    return hqm.get_or_create(question)
+
+
+# --- HITType ---
+
+
+@pytest.fixture
+def hit_type() -> HitType:
+    return HitType(
+        title="Awesome Surveys!",
+        description="Give us your opinion",
+        reward=USDCent(5),
+        keywords="market,research,amazing",
+        min_active=10,
+    )
+
+
+@pytest.fixture
+def hit_type_record(
+    pg_config: PostgresConfig, htm: "HitTypeManager", hit_type: HitType
+) -> Generator[HitType, None, None]:
+
+    hit_type.amt_hit_type_id = generate_amt_id()
+
+    ht = htm.get_or_create(hit_type)
+
+    yield ht
+
+    with pg_config.make_connection() as conn:
+        with conn.cursor() as c:
+            c.execute("DELETE FROM mtwerk_hittype WHERE id=%s", (ht.id,))
+            conn.commit()
+
+
+@pytest.fixture
+def hit_type_record_with_amt_id(
+    pg_config: PostgresConfig, htm: "HitTypeManager", hit_type: HitType
+) -> Generator[HitType, None, None]:
+    # This is a real hit type I've previously registered with amt (sandbox).
+    # It will always exist
+    hit_type.amt_hit_type_id = "3217B3DC4P5YW9DRV9R3X8O56V041J"
+
+    # Get or create our db
+    ht = htm.get_or_create(hit_type)
+
+    yield ht
+
+    with pg_config.make_connection() as conn:
+        with conn.cursor() as c:
+            c.execute("DELETE FROM mtwerk_hittype WHERE id=%s", (ht.id,))
+            conn.commit()
+
+
+# --- HIT ---
+
+
+@pytest.fixture
+def hit(
+    amt_hit_id: str, amt_hit_type_id: str, amt_group_id: str, question: HitQuestion
+) -> Hit:
+    now = datetime.now(tz=timezone.utc)
+
+    return Hit.model_validate(
+        dict(
+            amt_hit_id=amt_hit_id,
+            amt_hit_type_id=amt_hit_type_id,
+            amt_group_id=amt_group_id,
+            status=HitStatus.Assignable,
+            review_status=HitReviewStatus.NotReviewed,
+            creation_time=now,
+            expiration=now + timedelta(days=3),
+            hit_question_xml=question.xml,
+            qualification_requirements=[],
+            max_assignments=1,
+            assignment_pending_count=0,
+            assignment_available_count=1,
+            assignment_completed_count=0,
+            description="Description",
+            keywords="Keywords",
+            reward=USDCent(5),
+            title="Title",
+            question_id=question.id,
+            hit_type_id=None,
+        )
+    )
+
+
+@pytest.fixture
+def hit_record(
+    pg_config: PostgresConfig,
+    hm: "HitManager",
+    question_record: HitQuestion,
+    hit_type_record: HitType,
+    hit: Hit,
+    amt_hit_id: str,
+) -> Generator[Hit, None, None]:
+    """
+    Returns a hit that exists in our db, but does not in amazon (the amt ids
+        are random). The mtwerk_hittype and mtwerk_question records will also
+        exist (in the db)
+    """
+
+    hit.hit_type_id = hit_type_record.id
+    hit.amt_hit_id = amt_hit_id
+    hit.question_id = question_record.id
+
+    hm.create(hit)
+
+    yield hit
+
+    with pg_config.make_connection() as conn:
+        with conn.cursor() as c:
+            c.execute("DELETE FROM mtwerk_hit WHERE id=%s", (hit.id,))
+            conn.commit()
+
+
+@pytest.fixture
+def hit_in_amt(
+    hm: "HitManager", question_record: HitQuestion, hit_type_record_with_amt_id: HitType
+) -> Hit:
+    # Actually create a new HIT in amt (sandbox)
+    hit = AMTManager.create_hit_with_hit_type(
+        hit_type=hit_type_record_with_amt_id, question=question_record
+    )
+    # Create it in the DB
+    hm.create(hit)
+    return hit
+
+
+# --- Assignment ---
+
+
+@pytest.fixture
+def assignment_stub(
+    hit: Hit, amt_assignment_id: str, amt_worker_id: str
+) -> AssignmentStub:
+    now = datetime.now(tz=timezone.utc)
+    return AssignmentStub(
+        amt_assignment_id=amt_assignment_id,
+        amt_hit_id=hit.amt_hit_id,
+        amt_worker_id=amt_worker_id,
+        status=AssignmentStatus.Submitted,
+        modified_at=now,
+        created_at=now,
+    )
+
+
+@pytest.fixture
+def assignment_stub_record(
+    pg_config: PostgresConfig,
+    am: "AssignmentManager",
+    hit_record: Hit,
+    assignment_stub: AssignmentStub,
+) -> Generator[AssignmentStub, None, None]:
+    """
+    Returns an AssignmentStub that exists in our db, but does not in
+    amazon (the amt ids are random). The mtwerk_hit, mtwerk_hittype, and
+    mtwerk_question records will also exist (in the db)
+    """
+    assignment_stub.hit_id = hit_record.id
+    am.create_stub(stub=assignment_stub)
+
+    yield assignment_stub
+
+    with pg_config.make_connection() as conn:
+        with conn.cursor() as c:
+            c.execute(
+                "DELETE FROM mtwerk_assignment WHERE id=%s", (assignment_stub.id,)
+            )
+            conn.commit()
+
+
+@pytest.fixture
+def assignment(assignment_factory: Callable[..., Assignment]) -> Assignment:
+    return assignment_factory()
+
+
+@pytest.fixture
+def assignment_record(
+    pg_config: PostgresConfig,
+    hit_record: Hit,
+    assignment_record_factory: Callable[..., Assignment],
+) -> Generator[Assignment, None, None]:
+    assignment = assignment_record_factory(hit_id=hit_record.id)
+
+    yield assignment
+
+    with pg_config.make_connection() as conn:
+        with conn.cursor() as c:
+            c.execute(
+                "DELETE FROM mtwerk_assignment WHERE id=%s", (assignment_stub.id,)
+            )
+            conn.commit()
+
+
+@pytest.fixture
+def assignment_factory(hit: Hit) -> Callable[[Optional[str]], Assignment]:
+
+    def _inner(amt_worker_id: Optional[str] = None) -> Assignment:
+        now = datetime.now(tz=timezone.utc)
+        amt_assignment_id = generate_amt_id()
+        amt_worker_id = amt_worker_id or generate_amt_id()
+
+        return Assignment(
+            amt_assignment_id=amt_assignment_id,
+            amt_hit_id=hit.amt_hit_id,
+            amt_worker_id=amt_worker_id,
+            status=AssignmentStatus.Submitted,
+            modified_at=now,
+            created_at=now,
+            accept_time=now,
+            auto_approval_time=now,
+            submit_time=now,
+        )
+
+    return _inner
+
+
+@pytest.fixture
+def assignment_record_factory(
+    am: "AssignmentManager", assignment_factory: Callable[..., Assignment]
+) -> Callable[..., Assignment]:
+
+    def _inner(hit_id: int, amt_worker_id: Optional[str] = None) -> Assignment:
+        a = assignment_factory(amt_worker_id=amt_worker_id)
+        a.hit_id = hit_id
+        am.create_stub(a)
+        am.update_answer(a)
+        return a
+
+    return _inner