check_hit_status to use manager param, and amt_stub_context usage in flow tests.

2 files changed, 81 insertions, 35 deletions

diff --git a/tests/fixtures/flow.py b/tests/fixtures/flow.py
index 633e05d..5ee2026 100644
--- a/tests/fixtures/flow.py
+++ b/tests/fixtures/flow.py
@@ -26,8 +26,8 @@ from generalresearchutils.models.thl.definitions import PayoutStatus
 
 @pytest.fixture
 def approved_assignment_stubs(
-    get_assignment_response: GetAssignmentResponseTypeDef,
-    get_assignment_response_approved: Callable[[str], GetAssignmentResponseTypeDef],
+    assignment_response: GetAssignmentResponseTypeDef,
+    assignment_response_approved: Callable[[str], GetAssignmentResponseTypeDef],
     amt_assignment_id: str,
     amt_hit_id: str,
     get_hit_response_reviewing: GetHITResponseTypeDef,
@@ -42,9 +42,9 @@ def approved_assignment_stubs(
         override_approve_response: Optional[str] = None,
     ) -> list[Dict[str, Any]]:
 
-        response = override_response or get_assignment_response
-        approve_response = (
-            override_approve_response or get_assignment_response_approved(feedback)
+        response = override_response or assignment_response
+        approve_response = override_approve_response or assignment_response_approved(
+            feedback
         )
 
         return [
diff --git a/tests/flow/test_tasks.py b/tests/flow/test_tasks.py
index 82d4912..2aeffb9 100644
--- a/tests/flow/test_tasks.py
+++ b/tests/flow/test_tasks.py
@@ -4,7 +4,6 @@ from typing import Callable, Dict, Any
 
 import pytest
 from botocore.stub import Stubber
-from jb.decorators import AMT_CLIENT
 from jb.flow.assignment_tasks import process_assignment_submitted
 from jb.managers.amt import (
     AMTManager,
@@ -25,12 +24,14 @@ 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(responses):
+def amt_stub_context(amt_client: MTurkClient, responses: list[Dict[str, Any]]):
+
     # ty chatgpt for this
-    with Stubber(AMT_CLIENT) as stub:
+    with Stubber(amt_client) as stub:
         for r in responses:
             stub.add_response(
                 r["operation"],
@@ -44,6 +45,8 @@ class TestHITTasks:
 
     def test_fake_get_assignment(
         self,
+        amt_client: MTurkClient,
+        amtm: AMTManager,
         amt_assignment_id: str,
         amt_worker_id: str,
         assignment_response: GetAssignmentResponseTypeDef,
@@ -51,11 +54,12 @@ class TestHITTasks:
         # 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:
+
+        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(
+            assignment = amtm.get_assignment_if_exists(
                 amt_assignment_id=amt_assignment_id
             )
             assert assignment is not None
@@ -70,7 +74,11 @@ 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,
@@ -91,8 +99,12 @@ class TestProcessAssignmentSubmitted:
             reject_reason=REJECT_MESSAGE_UNKNOWN_ASSIGNMENT
         )
 
-        with amt_stub_context(amt_stubs) as stub, caplog.at_level(logging.WARNING):
-            process_assignment_submitted(mturk_event)
+        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
@@ -106,9 +118,11 @@ class TestProcessAssignmentSubmitted:
 
     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,
@@ -129,8 +143,12 @@ class TestProcessAssignmentSubmitted:
 
         mock_thl_responses(user_blocked=True)
 
-        with amt_stub_context(amt_stubs) as stub, caplog.at_level(logging.WARNING):
-            process_assignment_submitted(am=am, hm=hm, bm=bm, event=mturk_event)
+        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
@@ -143,8 +161,12 @@ class TestProcessAssignmentSubmitted:
 
     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,
@@ -171,8 +193,12 @@ class TestProcessAssignmentSubmitted:
 
         mock_thl_responses(user_blocked=False)
 
-        with amt_stub_context(amt_stubs) as stub, caplog.at_level(logging.WARNING):
-            process_assignment_submitted(mturk_event)
+        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
@@ -186,17 +212,22 @@ class TestProcessAssignmentSubmitted:
 
     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_rejected_no_tsid: GetAssignmentResponseTypeDef,
+        assignment_response_factory_rejected_no_tsid: Callable[
+            ..., GetAssignmentResponseTypeDef
+        ],
         assignment_response_no_tsid: GetAssignmentResponseTypeDef,
-        assignment_factory: Assignment,
-        assignment_stub_factory: Assignment,
+        assignment_factory: Callable[..., Assignment],
         hit_record: Hit,
         amt_worker_id: str,
     ):
@@ -212,22 +243,26 @@ class TestProcessAssignmentSubmitted:
         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_factory  # we need this to make the assignment stub in the db
+        _ = 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_rejected_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_stubs) as stub, caplog.at_level(logging.WARNING):
-            process_assignment_submitted(mturk_event)
+        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
@@ -241,17 +276,20 @@ class TestProcessAssignmentSubmitted:
 
     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_factory: Assignment,
+        assignment_stub_record: Assignment,
         caplog: pytest.LogCaptureFixture,
         mock_thl_responses: Callable[..., None],
         approved_assignment_stubs: Callable[..., list[Dict[str, Any]]],
     ):
-        _ = (
-            assignment_stub_factory()
-        )  # we need this to make the assignment stub in the db
+
+        _ = 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
@@ -261,10 +299,12 @@ class TestProcessAssignmentSubmitted:
         # 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)
+        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)
@@ -273,21 +313,27 @@ class TestProcessAssignmentSubmitted:
 
     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_factory: Assignment,
+        assignment_stub_record: Assignment,
         caplog: pytest.LogCaptureFixture,
         mock_thl_responses: Callable[..., None],
         approved_assignment_stubs_w_bonus: list[Dict[str, Any]],
     ):
-        _ = assignment_stub_factory  # we need this to make the assignment stub in the db
+        _ = 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(
-            approved_assignment_stubs_w_bonus
+            amt_client, approved_assignment_stubs_w_bonus
         ) as stub, caplog.at_level(logging.WARNING):
-            process_assignment_submitted(mturk_event)
+            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)