path: root/tests/fixtures/models.py

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
from psycopg.errors import ForeignKeyViolation

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

    try:
        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()

    except ForeignKeyViolation:
        pass  # DB gets dropped anyway, don't care


@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

    try:
        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()

    except ForeignKeyViolation:
        pass  # DB gets dropped anyway, don't care


# --- 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

    try:
        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()

    except ForeignKeyViolation:
        pass  # DB gets dropped anyway, don't care


@pytest.fixture
def hit_in_amt(
    amtm: AMTManager,
    hm: "HitManager",
    question_record: HitQuestion,
    hit_type_record_with_amt_id: HitType,
) -> Hit:
    # Actually create a new HIT in amt (sandbox)
    hit = amtm.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

    try:
        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()

    except ForeignKeyViolation:
        pass  # DB gets dropped anyway, don't care


@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

    try:
        with pg_config.make_connection() as conn:
            with conn.cursor() as c:
                c.execute("DELETE FROM mtwerk_assignment WHERE id=%s", (assignment.id,))
                conn.commit()

    except ForeignKeyViolation:
        pass  # DB gets dropped anyway, don't care


@pytest.fixture
def assignment_factory(hit_record: 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_record.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