path: root/tests/flow/test_tasks.py

import logging
from contextlib import contextmanager
from typing import Callable, Dict, Any

import pytest
from botocore.stub import Stubber
from jb.flow.assignment_tasks import process_assignment_submitted
from jb.managers.amt import (
    AMTManager,
    APPROVAL_MESSAGE,
    REJECT_MESSAGE_BADDIE,
    REJECT_MESSAGE_UNKNOWN_ASSIGNMENT,
    REJECT_MESSAGE_NO_WORK,
    NO_WORK_APPROVAL_MESSAGE,
)
from mypy_boto3_mturk.type_defs import (
    GetAssignmentResponseTypeDef,
)
from generalresearchutils.currency import USDCent
from jb.managers.assignment import AssignmentManager
from jb.managers.bonus import BonusManager
from jb.managers.hit import HitManager
from jb.models.definitions import AssignmentStatus
from jb.models.event import MTurkEvent
from jb.models.hit import Hit
from jb.models.assignment import Assignment, AssignmentStub
from mypy_boto3_mturk import MTurkClient


@contextmanager
def amt_stub_context(amt_client: MTurkClient, responses: list[Dict[str, Any]]):

    # ty chatgpt for this
    with Stubber(amt_client) as stub:
        for r in responses:
            stub.add_response(
                r["operation"],
                r.get("response", {}),
                r.get("expected_params", {}),
            )
        yield stub


class TestHITTasks:

    def test_fake_get_assignment(
        self,
        amt_client: MTurkClient,
        amtm: AMTManager,
        amt_assignment_id: str,
        amt_worker_id: str,
        assignment_response: GetAssignmentResponseTypeDef,
    ):
        # Testing just that this boto stubber works (we fake a response
        # using the real boto client)
        fake_response = assignment_response.copy()

        with Stubber(amt_client) as stub:
            expected_params = {"AssignmentId": amt_assignment_id}
            stub.add_response("get_assignment", fake_response, expected_params)

            assignment = amtm.get_assignment_if_exists(
                amt_assignment_id=amt_assignment_id
            )
            assert assignment is not None
            assert assignment.amt_assignment_id == amt_assignment_id
            assert assignment.amt_worker_id == amt_worker_id

            # Optionally, ensure all queued responses were used:
            stub.assert_no_pending_responses()


class TestProcessAssignmentSubmitted:

    def test_no_assignment_in_db(
        self,
        amtm: AMTManager,
        hm: HitManager,
        am: AssignmentManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        hit_record: Hit,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        caplog: pytest.LogCaptureFixture,
        rejected_assignment_stubs: Callable[..., list[Dict[str, Any]]],
    ):

        # These records are auto cleaned up, so we need to explicitly create
        # a HIT record in the DB so the process_assignment_submitted task
        # doesn't error when we try to process the Request
        _ = hit_record

        # An assignment is submitted. The hit exists in the DB. The amt
        # assignment id is valid, but the assignment stub is not in our
        # db. Reject it and write the assignment to the db.

        amt_stubs = rejected_assignment_stubs(
            reject_reason=REJECT_MESSAGE_UNKNOWN_ASSIGNMENT
        )

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" in caplog.text
        assert f"Rejected assignment doesn't exist in DB. Creating ... " in caplog.text
        assert f"Rejected assignment: " in caplog.text
        stub.assert_no_pending_responses()

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Rejected
        assert ass.requester_feedback == REJECT_MESSAGE_UNKNOWN_ASSIGNMENT

    def test_assignment_in_db_user_doesnt_exist(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        assignment_stub_record: AssignmentStub,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        rejected_assignment_stubs: Callable[..., list[Dict[str, Any]]],
    ):
        # An assignment is submitted. The hit and AssignmentStub exist in the
        # DB. We think we're going to approve the Assignment, but the
        # user-profile / check blocked call on THL shows the user doesn't
        # exist (same thing would happen if the user does exist and is
        # blocked). So we reject.

        # We need this to make the assignment stub in the db
        _ = assignment_stub_record

        amt_stubs = rejected_assignment_stubs(reject_reason=REJECT_MESSAGE_BADDIE)

        mock_thl_responses(user_blocked=True)

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, am=am, hm=hm, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" not in caplog.text
        assert f"blocked or not exists" in caplog.text
        assert f"Rejected assignment: " in caplog.text

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Rejected
        assert ass.requester_feedback == REJECT_MESSAGE_BADDIE

    def test_no_work_w_warning(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        mturk_event: MTurkEvent,
        amt_client: MTurkClient,
        amt_assignment_id: str,
        assignment_stub_record: AssignmentStub,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        approved_assignment_stubs: Callable[..., list[Dict[str, Any]]],
        assignment_response_approved_no_tsid: GetAssignmentResponseTypeDef,
        assignment_response_no_tsid: GetAssignmentResponseTypeDef,
    ):
        # An Assignment is submitted. The hit and AssignmentStub exist in
        # the DB. The assignment has no tsid.
        # We APPROVE this assignment b/c we are very nice and give users a
        #   couple chances, with an explanation, before rejecting.

        # We need this to make the assignment stub in the db
        _ = assignment_stub_record

        # Simulate that the AMT.get_assignment call returns the assignment,
        #   but the answers XML has no tsid.
        amt_stubs = approved_assignment_stubs(
            feedback=NO_WORK_APPROVAL_MESSAGE,
            override_response=assignment_response_no_tsid,
            override_approve_response=assignment_response_approved_no_tsid,
        )

        mock_thl_responses(user_blocked=False)

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" not in caplog.text
        assert f"Assignment submitted with no tsid" in caplog.text
        assert f"Approved assignment: " in caplog.text

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Approved
        assert ass.requester_feedback == NO_WORK_APPROVAL_MESSAGE
        assert am.missing_tsid_count(amt_worker_id=ass.amt_worker_id) == 1

    def test_no_work_no_warning(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        mturk_event: MTurkEvent,
        amt_client: MTurkClient,
        amt_assignment_id: str,
        assignment_stub_record: AssignmentStub,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        rejected_assignment_stubs: Callable[..., list[Dict[str, Any]]],
        assignment_response_factory_rejected_no_tsid: Callable[
            ..., GetAssignmentResponseTypeDef
        ],
        assignment_response_no_tsid: GetAssignmentResponseTypeDef,
        assignment_factory: Callable[..., Assignment],
        hit_record: Hit,
        amt_worker_id: str,
    ):
        # An assignment is submitted. The hit and assignment stub exist in the DB.
        # The assignment has no tsid.

        # Going to create and submit 3 assignments w no work
        #  (all on the same hit, which we don't do in JB for real,
        #   but doesn't matter here)
        _a1 = assignment_factory(hit_id=hit_record.id, amt_worker_id=amt_worker_id)
        _a2 = assignment_factory(hit_id=hit_record.id, amt_worker_id=amt_worker_id)
        _a3 = assignment_factory(hit_id=hit_record.id, amt_worker_id=amt_worker_id)
        assert am.missing_tsid_count(amt_worker_id=amt_worker_id) == 3
        # So now, we'll reject, b/c they've already gotten 3 warnings

        _ = assignment_stub_record  # we need this to make the assignment stub in the db

        # Simulate that the AMT.get_assignment call returns the assignment, but the answers xml
        #   has no tsid.
        amt_stubs = rejected_assignment_stubs(
            reject_reason=REJECT_MESSAGE_NO_WORK,
            override_response=assignment_response_no_tsid,
            override_reject_response=assignment_response_factory_rejected_no_tsid(
                REJECT_MESSAGE_NO_WORK
            ),
        )

        mock_thl_responses(user_blocked=False)

        with amt_stub_context(amt_client, amt_stubs) as stub, caplog.at_level(
            logging.WARNING
        ):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        assert f"No assignment found in DB: {amt_assignment_id}" not in caplog.text
        assert f"Assignment submitted with no tsid" in caplog.text
        assert f"Rejected assignment: " in caplog.text

        # It will exist in the db since we can validate the model.
        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Rejected
        assert ass.requester_feedback == REJECT_MESSAGE_NO_WORK

    def test_assignment_submitted_no_bonus(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        assignment_stub_record: Assignment,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        approved_assignment_stubs: Callable[..., list[Dict[str, Any]]],
    ):

        _ = assignment_stub_record  # we need this to make the assignment stub in the db

        # The "send bonus" stuff will still run, even if the user didn't get
        # a complete, because all we do is check the user's wallet balance (if
        # an assignment is approved) and they may have money in their wallet
        # from a prev event or bribe
        #
        # So mock the wallet balance as 1cent, so no bonus will be triggered

        mock_thl_responses(status_complete=False, wallet_redeemable_amount=1)
        with amt_stub_context(
            amt_client, approved_assignment_stubs()
        ) as stub, caplog.at_level(logging.WARNING):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Approved
        assert ass.requester_feedback == APPROVAL_MESSAGE

    def test_assignment_submitted_w_bonus(
        self,
        amtm: AMTManager,
        am: AssignmentManager,
        hm: HitManager,
        bm: BonusManager,
        amt_client: MTurkClient,
        mturk_event: MTurkEvent,
        amt_assignment_id: str,
        assignment_stub_record: Assignment,
        caplog: pytest.LogCaptureFixture,
        mock_thl_responses: Callable[..., None],
        approved_assignment_stubs_w_bonus: list[Dict[str, Any]],
    ):
        _ = assignment_stub_record  # we need this to make the assignment stub in the db
        mock_thl_responses(status_complete=True, wallet_redeemable_amount=10)

        with amt_stub_context(
            amt_client, approved_assignment_stubs_w_bonus
        ) as stub, caplog.at_level(logging.WARNING):
            process_assignment_submitted(
                amtm=amtm, hm=hm, am=am, bm=bm, event=mturk_event
            )
            stub.assert_no_pending_responses()

        ass = am.get(amt_assignment_id=amt_assignment_id)
        assert ass.status == AssignmentStatus.Approved
        assert ass.requester_feedback == APPROVAL_MESSAGE

        bonus = bm.filter(amt_assignment_id=amt_assignment_id)[0]
        assert bonus.amount == USDCent(7)