1 files changed, 43 insertions, 9 deletions

diff --git a/jb/flow/assignment_tasks.py b/jb/flow/assignment_tasks.py
index 4022716..bb0877d 100644
--- a/jb/flow/assignment_tasks.py
+++ b/jb/flow/assignment_tasks.py
@@ -5,6 +5,7 @@ from typing import Optional
 
 from generalresearchutils.models.thl.definitions import PayoutStatus, StatusCode1
 from generalresearchutils.models.thl.wallet.cashout_method import CashoutRequestInfo
+from generalresearchutils.currency import USDCent
 
 from jb.decorators import AM, HM, BM
 from jb.flow.monitoring import emit_error_event, emit_assignment_event, emit_bonus_event
@@ -21,21 +22,20 @@ from jb.managers.thl import (
     get_user_blocked,
     get_task_status,
     user_cashout_request,
-    AMT_ASSIGNMENT_CASHOUT_METHOD,
     manage_pending_cashout,
     get_user_blocked_or_not_exists,
     get_wallet_balance,
-    AMT_BONUS_CASHOUT_METHOD,
 )
 from jb.models.assignment import Assignment
-from jb.models.currency import USDCent
 from jb.models.definitions import AssignmentStatus
 from jb.models.event import MTurkEvent
+from jb.config import settings
 
 
 def process_assignment_submitted(event: MTurkEvent) -> None:
     """
-    Called either directly or from the SNS Notification that a HIT was submitted
+    Called either directly or from the SNS Notification that a
+    HIT was submitted
 
     :return: None
     """
@@ -137,13 +137,22 @@ def process_assignment_submitted(event: MTurkEvent) -> None:
         return issue_worker_payment(assignment)
 
 
-def review_hit(assignment):
+def review_hit(assignment: Assignment) -> None:
     # Reviewable to Reviewing
     AMTManager.update_hit_review_status(amt_hit_id=assignment.amt_hit_id, revert=False)
     hit, _ = AMTManager.get_hit_if_exists(amt_hit_id=assignment.amt_hit_id)
+
+    if hit is None:
+        logging.warning(
+            f"Hit not found when trying to review hit: {assignment.amt_hit_id}"
+        )
+        return None
+
     # Update the db
     HM.update_hit(hit)
 
+    return None
+
 
 def handle_assignment_w_no_work(assignment: Assignment) -> Assignment:
     """
@@ -267,6 +276,10 @@ def handle_assignment_w_work(assignment: Assignment) -> Assignment:
     amt_worker_id = assignment.amt_worker_id
     amt_assignment_id = assignment.amt_assignment_id
     tsid = assignment.tsid
+    assert (
+        tsid is not None
+    ), "Assignment must have a tsid to be handled in handle_assignment_w_work"
+
     hit = HM.get_from_amt_id(amt_hit_id=assignment.amt_hit_id)
 
     tsr = get_task_status(tsid=tsid)
@@ -289,6 +302,7 @@ def handle_assignment_w_work(assignment: Assignment) -> Assignment:
             event_type = "assignment_submitted_quality_fail"
         else:
             event_type = "assignment_submitted_work_not_complete"
+
         emit_error_event(
             event_type=event_type,
             amt_hit_type_id=hit.amt_hit_type_id,
@@ -324,6 +338,8 @@ def handle_assignment_w_work(assignment: Assignment) -> Assignment:
         )
         return assignment
 
+    assert req.id
+
     # We've approved the HIT payment, now update the db to reflect this, and approve the assignment
     assignment = approve_assignment(
         amt_assignment_id=amt_assignment_id,
@@ -331,7 +347,9 @@ def handle_assignment_w_work(assignment: Assignment) -> Assignment:
         amt_hit_type_id=hit.amt_hit_type_id,
     )
     # We complete after the assignment is approved
-    complete_res = manage_pending_cashout(req.id, PayoutStatus.COMPLETE)
+    complete_res = manage_pending_cashout(
+        cashout_id=req.id, payout_status=PayoutStatus.COMPLETE
+    )
     if complete_res.status != PayoutStatus.COMPLETE:
         # unclear wny this would happen
         raise ValueError(f"Failed to complete cashout: {req.id}")
@@ -345,8 +363,9 @@ def submit_and_approve_amt_assignment_request(
     req = user_cashout_request(
         amt_worker_id=amt_worker_id,
         amount=amount,
-        cashout_method_id=AMT_ASSIGNMENT_CASHOUT_METHOD,
+        cashout_method_id=settings.amt_assignment_cashout_method,
     )
+    assert req.id
 
     if req.status != PayoutStatus.PENDING:
         return None
@@ -365,8 +384,9 @@ def submit_and_approve_amt_bonus_request(
     req = user_cashout_request(
         amt_worker_id=amt_worker_id,
         amount=amount,
-        cashout_method_id=AMT_BONUS_CASHOUT_METHOD,
+        cashout_method_id=settings.amt_bonus_cashout_method,
     )
+    assert req.id
 
     if req.status != PayoutStatus.PENDING:
         return None
@@ -404,6 +424,7 @@ def issue_worker_payment(assignment: Assignment) -> None:
             amt_hit_type_id=hit.amt_hit_type_id,
         )
         return None
+    assert pe.id
 
     AMTManager.send_bonus(
         amt_worker_id=assignment.amt_worker_id,
@@ -412,13 +433,26 @@ def issue_worker_payment(assignment: Assignment) -> None:
         reason=BONUS_MESSAGE,
         unique_request_token=pe.id,
     )
+
     # Confirm it was sent through amt
     bonus = AMTManager.get_bonus(
         amt_assignment_id=assignment.amt_assignment_id, payout_event_id=pe.id
     )
+
+    if bonus is None:
+        logging.warning(
+            f"Failed to find bonus after sending it: {amt_assignment_id} {pe.id}"
+        )
+        emit_error_event(
+            event_type="bonus_not_found_after_sending",
+            amt_hit_type_id=hit.amt_hit_type_id,
+        )
+        return None
+
     # Create in DB
-    BM.create(bonus)
+    BM.create(bonus=bonus)
     emit_bonus_event(amount=amount, amt_hit_type_id=hit.amt_hit_type_id)
+
     # Complete cashout
     res = manage_pending_cashout(pe.id, PayoutStatus.COMPLETE)
     if res.status != PayoutStatus.COMPLETE: