Extensive use of type checking. Movement of pytest conf towards handling managers (for db agnostic unittest). Starting to organize pytests.

1 files changed, 25 insertions, 15 deletions

diff --git a/jb/managers/thl.py b/jb/managers/thl.py
index b1dcbde..83f49f6 100644
--- a/jb/managers/thl.py
+++ b/jb/managers/thl.py
@@ -1,7 +1,3 @@
-from decimal import Decimal
-from typing import Dict, Optional
-
-import requests
 from generalresearchutils.models.thl.payout import UserPayoutEvent
 from generalresearchutils.models.thl.task_status import TaskStatusResponse
 from generalresearchutils.models.thl.wallet.cashout_method import (
@@ -9,45 +5,58 @@ from generalresearchutils.models.thl.wallet.cashout_method import (
     CashoutRequestInfo,
 )
 
+from generalresearchutils.models.thl.user_profile import UserProfile
+from generalresearchutils.currency import USDCent
+
 from jb.config import settings
-from jb.models.currency import USDCent
-from jb.models.definitions import PayoutStatus
 
-# TODO: Organize this more with other endpoints (offerwall, cashout requests/approvals, etc).
+from generalresearchutils.models.thl.definitions import PayoutStatus
+
+
+from typing import Optional
+import requests
+
+# TODO: Organize this more with other endpoints (offerwall, cashout
+# requests/approvals, etc).
 
 
-def get_user_profile(amt_worker_id: str) -> Dict:
+def get_user_profile(amt_worker_id: str) -> UserProfile:
     url = f"{settings.fsb_host}{settings.product_id}/user/{amt_worker_id}/profile/"
     res = requests.get(url).json()
     if res.get("detail") == "user not found":
         raise ValueError("user not found")
-    return res["user_profile"]
+
+    return UserProfile.model_validate(res["user_profile"])
 
 
 def get_user_blocked(amt_worker_id: str) -> bool:
+    # Not blocked if None
     res = get_user_profile(amt_worker_id=amt_worker_id)
-    return res["user"]["blocked"]
+    return res.user.blocked if res.user.blocked is not None else False
 
 
-def get_user_blocked_or_not_exists(amt_worker_id: str) -> bool:
+def get_user_blocked_or_not_exists(amt_worker_id: str) -> Optional[bool]:
     try:
         res = get_user_profile(amt_worker_id=amt_worker_id)
-        return res["user"]["blocked"]
+        return res.user.blocked if res.user.blocked is not None else False
     except ValueError as e:
         if e.args[0] == "user not found":
             return True
 
+        return None
+
 
 def get_task_status(tsid: str) -> Optional[TaskStatusResponse]:
     url = f"{settings.fsb_host}{settings.product_id}/status/{tsid}/"
     d = requests.get(url).json()
     if d.get("msg") == "invalid tsid":
         return None
+
     return TaskStatusResponse.model_validate(d)
 
 
 def user_cashout_request(
-    amt_worker_id: str, amount: USDCent, cashout_method_id
+    amt_worker_id: str, amount: USDCent, cashout_method_id: str
 ) -> CashoutRequestInfo:
     assert cashout_method_id in {
         settings.amt_assignment_cashout_method,
@@ -56,7 +65,8 @@ def user_cashout_request(
     assert isinstance(amount, USDCent)
     assert USDCent(0) < amount < USDCent(10_00)
     url = f"{settings.fsb_host}{settings.product_id}/cashout/"
-    body = {
+
+    body: dict[str, str | int] = {
         "bpuid": amt_worker_id,
         "amount": int(amount),
         "cashout_method_id": cashout_method_id,
@@ -81,7 +91,7 @@ def manage_pending_cashout(
     return UserPayoutEvent.model_validate(d)
 
 
-def get_wallet_balance(amt_worker_id: str):
+def get_wallet_balance(amt_worker_id: str) -> USDCent:
     url = f"{settings.fsb_host}{settings.product_id}/wallet/"
     params = {"bpuid": amt_worker_id}
     return USDCent(requests.get(url, params=params).json()["wallet"]["amount"])