get_wallet_balance_if_non_negative

1 files changed, 11 insertions, 1 deletions

diff --git a/jb/managers/thl.py b/jb/managers/thl.py
index c1cd672..6e8effc 100644
--- a/jb/managers/thl.py
+++ b/jb/managers/thl.py
@@ -35,7 +35,7 @@ def get_user_profile(amt_worker_id: str) -> UserProfile:
     # todo: this contains computed fields inside each streak object
     user_profile.pop("streaks", None)
     # todo: this shouldn't be in here anyways ---v
-    user_profile['user'].pop("id", None)
+    user_profile["user"].pop("id", None)
 
     return UserProfile.model_validate(user_profile)
 
@@ -104,6 +104,16 @@ def manage_pending_cashout(
 
 
 def get_wallet_balance(amt_worker_id: str) -> USDCent:
+    # This will raise an Exception if wallet balance is negative
     url = f"{settings.fsb_host}{settings.product_id}/wallet/"
     params = {"bpuid": amt_worker_id}
     return USDCent(requests.get(url, params=params).json()["wallet"]["amount"])
+
+
+def get_wallet_balance_if_non_negative(amt_worker_id: str) -> Optional[USDCent]:
+    url = f"{settings.fsb_host}{settings.product_id}/wallet/"
+    params = {"bpuid": amt_worker_id}
+    amt = requests.get(url, params=params).json()["wallet"]["amount"]
+    if amt >= 0:
+        return USDCent(amt)
+    return None