1 files changed, 30 insertions, 9 deletions
diff --git a/jb/api/magic_token.py b/jb/api/magic_token.py
index 7b6c1aa..910dc12 100644
--- a/jb/api/magic_token.py
+++ b/jb/api/magic_token.py
@@ -2,8 +2,9 @@ import hashlib
 import secrets
 
 from fastapi import HTTPException, status
+from generalresearch.redis_helper import RedisConfig
 
-from jb.decorators import get_redis
+from jb.decorators import get_redis_config
 from jb.models.auth import AmtAccountLink
 
 MAGIC_TOKEN_PREFIX = "auth:magic:"
@@ -17,7 +18,7 @@ def redis_token_key(token: str, prefix: str = MAGIC_TOKEN_PREFIX) -> str:
     return f"{prefix}{digest}"
 
 
-def create_magic_token(user_email: str) -> str:
+def create_magic_token(user_email: str, redis_config: RedisConfig | None = None) -> str:
     """Create a short-lived, single-use token for a user.
     The raw token can then be sent by email.
     """
@@ -25,7 +26,12 @@ def create_magic_token(user_email: str) -> str:
         raise ValueError("user_email must not be empty")
 
     token = secrets.token_urlsafe(32)
-    redis_client = get_redis()
+
+    if redis_config is None:
+        redis_config = get_redis_config()
+
+    redis_client = redis_config.create_redis_client()
+
     redis_client.set(
         redis_token_key(token),
         user_email,
@@ -34,8 +40,12 @@ def create_magic_token(user_email: str) -> str:
     return token
 
 
-def consume_magic_token(token: str) -> str:
-    redis_client = get_redis()
+def consume_magic_token(token: str, redis_config: RedisConfig | None = None) -> str:
+    if redis_config is None:
+        redis_config = get_redis_config()
+
+    redis_client = redis_config.create_redis_client()
+
     user_email = redis_client.getdel(redis_token_key(token))
     if user_email is None:
         raise HTTPException(
@@ -45,9 +55,15 @@ def consume_magic_token(token: str) -> str:
     return str(user_email)
 
 
-def create_amt_account_link_token(email: str, amt_worker_id: str) -> str:
+def create_amt_account_link_token(
+    email: str, amt_worker_id: str, redis_config: RedisConfig | None = None
+) -> str:
     """Bind an email and AMT worker ID to an opaque, short-lived token."""
-    redis_client = get_redis()
+    if redis_config is None:
+        redis_config = get_redis_config()
+
+    redis_client = redis_config.create_redis_client()
+
     data = AmtAccountLink(
         email=email,
         amt_worker_id=amt_worker_id,
@@ -61,9 +77,14 @@ def create_amt_account_link_token(email: str, amt_worker_id: str) -> str:
     return token
 
 
-def consume_amt_account_link_token(token: str) -> AmtAccountLink:
+def consume_amt_account_link_token(
+    token: str, redis_config: RedisConfig | None = None
+) -> AmtAccountLink:
     """Atomically consume and validate an AMT account-link token."""
-    redis_client = get_redis()
+    if redis_config is None:
+        redis_config = get_redis_config()
+    redis_client = redis_config.create_redis_client()
+
     raw_data = redis_client.getdel(
         redis_token_key(token, AMT_ACCOUNT_LINK_TOKEN_PREFIX)
     )