import logging from contextlib import contextmanager from datetime import timezone, datetime from typing import Dict from uuid import uuid4 import pytest import requests from botocore.stub import Stubber from generalresearchutils.models.thl.payout import UserPayoutEvent from generalresearchutils.models.thl.wallet import PayoutType from generalresearchutils.models.thl.wallet.cashout_method import ( CashoutRequestResponse, CashoutRequestInfo, ) from jb.config import settings from jb.decorators import AMT_CLIENT, AM, BM 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, BONUS_MESSAGE, NO_WORK_APPROVAL_MESSAGE, ) from jb.models.currency import USDCent from jb.models.definitions import AssignmentStatus, PayoutStatus from jb.models.event import MTurkEvent @contextmanager def amt_stub_context(responses): # 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 @pytest.fixture def approved_assignment_stubs( get_assignment_response, get_assignment_response_approved, amt_assignment_id, amt_hit_id, get_hit_response_reviewing, ): # These are the AMT_CLIENT stubs/mocks that need to be set when running # process_assignment_submitted() which will result in an approved # assignment and sent bonus def _approved_assignment_stubs( feedback: str = APPROVAL_MESSAGE, override_response=None, override_approve_response=None, ): response = override_response or get_assignment_response approve_response = ( override_approve_response or get_assignment_response_approved(feedback) ) return [ { "operation": "get_assignment", "response": response, "expected_params": {"AssignmentId": amt_assignment_id}, }, { "operation": "approve_assignment", "response": {}, "expected_params": { "AssignmentId": amt_assignment_id, "RequesterFeedback": feedback, "OverrideRejection": False, }, }, { "operation": "get_assignment", "response": approve_response, "expected_params": {"AssignmentId": amt_assignment_id}, }, { "operation": "update_hit_review_status", "response": {}, "expected_params": {"HITId": amt_hit_id, "Revert": False}, }, { "operation": "get_hit", "response": get_hit_response_reviewing, "expected_params": {"HITId": amt_hit_id}, }, ] return _approved_assignment_stubs @pytest.fixture def approved_assignment_stubs_w_bonus( approved_assignment_stubs, amt_worker_id, amt_assignment_id, pe_id ): now = datetime.now(tz=timezone.utc) stubs = approved_assignment_stubs().copy() stubs.append( { "operation": "send_bonus", "response": {}, "expected_params": { "WorkerId": amt_worker_id, "BonusAmount": "0.07", "AssignmentId": amt_assignment_id, "Reason": BONUS_MESSAGE, "UniqueRequestToken": pe_id, }, } ) stubs.append( { "operation": "list_bonus_payments", "response": { "BonusPayments": [ { "WorkerId": amt_worker_id, "BonusAmount": "0.07", "AssignmentId": amt_assignment_id, "Reason": BONUS_MESSAGE, "GrantTime": now, } ] }, "expected_params": {"AssignmentId": amt_assignment_id}, } ) return stubs @pytest.fixture def rejected_assignment_stubs( get_assignment_response, get_assignment_response_rejected, amt_assignment_id, amt_hit_id, get_hit_response_reviewing, ): # These are the AMT_CLIENT stubs/mocks that need to be set when running # process_assignment_submitted() which will result in a rejected # assignment def _rejected_assignment_stubs( reject_reason: str, override_response=None, override_reject_response=None ): response = override_response or get_assignment_response reject_response = override_reject_response or get_assignment_response_rejected( reject_reason ) return [ { "operation": "get_assignment", "response": response, "expected_params": {"AssignmentId": amt_assignment_id}, }, { "operation": "reject_assignment", "response": {}, "expected_params": { "AssignmentId": amt_assignment_id, "RequesterFeedback": reject_reason, }, }, { "operation": "get_assignment", "response": reject_response, "expected_params": {"AssignmentId": amt_assignment_id}, }, { "operation": "update_hit_review_status", "response": {}, "expected_params": {"HITId": amt_hit_id, "Revert": False}, }, { "operation": "get_hit", "response": get_hit_response_reviewing, "expected_params": {"HITId": amt_hit_id}, }, ] return _rejected_assignment_stubs @pytest.fixture def mock_thl_responses(monkeypatch, amt_worker_id, tsid, pe_id): original_get = requests.get original_post = requests.post class MockThlCashoutRequestResponse: def json(self): return CashoutRequestResponse( status="success", cashout=CashoutRequestInfo( id=pe_id, description="amt something", status=PayoutStatus.PENDING, ), ).model_dump(mode="json") def _apply_mock( user_blocked: bool = False, status_finished: bool = True, status_complete: bool = False, wallet_redeemable_amount: int = 10, ): def mock_get(url, *args, **kwargs): profile_url = f"{settings.fsb_host}{settings.product_id}/user/{amt_worker_id}/profile/" status_url = f"{settings.fsb_host}{settings.product_id}/status/{tsid}/" cashout_request_url = f"{settings.fsb_host}{settings.product_id}/cashout/" wallet_url = f"{settings.fsb_host}{settings.product_id}/wallet/" if url == profile_url: class MockThlProfileResponse: def json(self): return {"user-profile": {"user": {"blocked": user_blocked}}} return MockThlProfileResponse() elif url == wallet_url: class MockThlWalletResponse: def json(self): return { "wallet": {"redeemable_amount": wallet_redeemable_amount} } return MockThlWalletResponse() elif url == status_url: class MockThlStatusResponse: def json(self): return { "tsid": tsid, "product_id": str(settings.product_id), "product_user_id": amt_worker_id, "started": "2020-06-02T00:30:35.036398Z", "finished": ( "2020-06-02T00:31:35.036398Z" if status_finished else None ), "status": 3 if status_complete else 2, "payout": 10 if status_complete else 0, "user_payout": 10 if status_complete else 0, "status_code_1": "BUYER_FAIL", "status_code_2": None, } return MockThlStatusResponse() elif url == cashout_request_url: return MockThlCashoutRequestResponse() else: raise ValueError(f"unhandled call: {url=} {args=} {kwargs=}") return original_get(url, *args, **kwargs) def mock_post(url, *args, **kwargs): cashout_request_url = f"{settings.fsb_host}{settings.product_id}/cashout/" manage_cashout_request_url = f"{settings.fsb_host}{settings.fsb_host_private_route}/thl/manage_cashout/" if url == cashout_request_url: return MockThlCashoutRequestResponse() elif url == manage_cashout_request_url: json = kwargs["json"] print(json) payout_id = json["payout_id"] new_status = json["new_status"] class MockThlManageCashoutResponse: def json(self): return UserPayoutEvent( uuid=payout_id, status=new_status, amount=USDCent(5), debit_account_uuid=uuid4().hex, cashout_method_uuid=uuid4().hex, payout_type=PayoutType.AMT, ).model_dump(mode="json") return MockThlManageCashoutResponse() return original_post(url, *args, **kwargs) monkeypatch.setattr(requests, "get", mock_get) monkeypatch.setattr(requests, "post", mock_post) return _apply_mock @pytest.fixture def mturk_event(amt_assignment_id, amt_hit_id, amt_hit_type_id): 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, ) def test_fake_get_assignment( amt_assignment_id, amt_worker_id, get_assignment_response: Dict ): # Testing just that this boto stubber works (we fake a response using the real boto client) fake_response = get_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 = AMTManager.get_assignment_if_exists( amt_assignment_id=amt_assignment_id ) 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, mturk_event, amt_assignment_id, get_assignment_response: Dict, caplog, hit_in_db, rejected_assignment_stubs, ): # 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_stubs) as stub, caplog.at_level(logging.WARNING): process_assignment_submitted(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, mturk_event, amt_assignment_id, assignment_stub_in_db, caplog, mock_thl_responses, rejected_assignment_stubs, ): # An assignment is submitted. The hit and assignment stub 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. _ = assignment_stub_in_db # we need this to make the assignment stub in the db amt_stubs = rejected_assignment_stubs(reject_reason=REJECT_MESSAGE_BADDIE) mock_thl_responses(user_blocked=True) with amt_stub_context(amt_stubs) as stub, caplog.at_level(logging.WARNING): process_assignment_submitted(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, mturk_event, amt_assignment_id, assignment_stub_in_db, caplog, mock_thl_responses, approved_assignment_stubs, get_assignment_response_approved_no_tsid, get_assignment_response_no_tsid, ): # An assignment is submitted. The hit and assignment stub 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. _ = assignment_stub_in_db # 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 = approved_assignment_stubs( feedback=NO_WORK_APPROVAL_MESSAGE, override_response=get_assignment_response_no_tsid, override_approve_response=get_assignment_response_approved_no_tsid, ) mock_thl_responses(user_blocked=False) with amt_stub_context(amt_stubs) as stub, caplog.at_level(logging.WARNING): process_assignment_submitted(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, mturk_event, amt_assignment_id, assignment_stub_in_db, caplog, mock_thl_responses, rejected_assignment_stubs, get_assignment_response_rejected_no_tsid, get_assignment_response_no_tsid, assignment_in_db_factory, hit_in_db, amt_worker_id, ): # 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_in_db_factory(hit_id=hit_in_db.id, amt_worker_id=amt_worker_id) a2 = assignment_in_db_factory(hit_id=hit_in_db.id, amt_worker_id=amt_worker_id) a3 = assignment_in_db_factory(hit_id=hit_in_db.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_in_db # 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=get_assignment_response_no_tsid, override_reject_response=get_assignment_response_rejected_no_tsid( REJECT_MESSAGE_NO_WORK ), ) mock_thl_responses(user_blocked=False) with amt_stub_context(amt_stubs) as stub, caplog.at_level(logging.WARNING): process_assignment_submitted(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, mturk_event, amt_assignment_id, assignment_stub_in_db, caplog, mock_thl_responses, approved_assignment_stubs, ): _ = assignment_stub_in_db # 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(approved_assignment_stubs()) as stub, caplog.at_level( logging.WARNING ): process_assignment_submitted(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, mturk_event, amt_assignment_id, assignment_stub_in_db, caplog, mock_thl_responses, approved_assignment_stubs_w_bonus, ): _ = assignment_stub_in_db # 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( approved_assignment_stubs_w_bonus ) as stub, caplog.at_level(logging.WARNING): process_assignment_submitted(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)