From 37b41d42d58b8e48f65db5303c1114bf27f58c0c Mon Sep 17 00:00:00 2001
From: rishabhpoddar <rishabh.poddar@gmail.com>
Date: Tue, 1 Oct 2024 16:52:40 +0530
Subject: [PATCH] fixes more tests

---
 dev-requirements.txt                          |  1 +
 supertokens_python/auth_utils.py              | 12 ++++----
 .../multi_factor_auth_claim.py                |  5 +++-
 .../recipe/totp/recipe_implementation.py      |  3 +-
 tests/auth-react/flask-server/app.py          | 29 ++++++++++++++++++-
 5 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/dev-requirements.txt b/dev-requirements.txt
index 4291ad283..0694d481a 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -85,3 +85,4 @@ uvicorn==0.18.2
 Werkzeug==2.0.3
 wrapt==1.13.3
 zipp==3.7.0
+pyotp==2.9.0
\ No newline at end of file
diff --git a/supertokens_python/auth_utils.py b/supertokens_python/auth_utils.py
index 7e1bcff7c..e471ea522 100644
--- a/supertokens_python/auth_utils.py
+++ b/supertokens_python/auth_utils.py
@@ -452,9 +452,9 @@ class OkFirstFactorResponse:
 
 
 class OkSecondFactorLinkedResponse:
-    status: Literal["OK"]
-    is_first_factor: Literal[False]
-    input_user_already_linked_to_session_user: Literal[True]
+    status: Literal["OK"] = "OK"
+    is_first_factor: Literal[False] = False
+    input_user_already_linked_to_session_user: Literal[True] = True
     session_user: User
 
     def __init__(self, session_user: User):
@@ -462,9 +462,9 @@ def __init__(self, session_user: User):
 
 
 class OkSecondFactorNotLinkedResponse:
-    status: Literal["OK"]
-    is_first_factor: Literal[False]
-    input_user_already_linked_to_session_user: Literal[False]
+    status: Literal["OK"] = "OK"
+    is_first_factor: Literal[False] = False
+    input_user_already_linked_to_session_user: Literal[False] = False
     session_user: User
     linking_to_session_user_requires_verification: bool
 
diff --git a/supertokens_python/recipe/multifactorauth/multi_factor_auth_claim.py b/supertokens_python/recipe/multifactorauth/multi_factor_auth_claim.py
index 3c4bd6551..8b45108ba 100644
--- a/supertokens_python/recipe/multifactorauth/multi_factor_auth_claim.py
+++ b/supertokens_python/recipe/multifactorauth/multi_factor_auth_claim.py
@@ -259,7 +259,10 @@ def remove_from_payload_by_merge_(
     def get_value_from_payload(
         self, payload: JSONObject, user_context: Optional[Dict[str, Any]] = None
     ) -> Optional[MFAClaimValue]:
-        return payload.get(self.key)
+        value = payload.get(self.key)
+        if value is None:
+            return None
+        return MFAClaimValue(c=value["c"], v=value["v"])
 
 
 MultiFactorAuthClaim = MultiFactorAuthClaimClass()
diff --git a/supertokens_python/recipe/totp/recipe_implementation.py b/supertokens_python/recipe/totp/recipe_implementation.py
index b213c4eb2..ce213dd88 100644
--- a/supertokens_python/recipe/totp/recipe_implementation.py
+++ b/supertokens_python/recipe/totp/recipe_implementation.py
@@ -107,10 +107,11 @@ async def create_device(
 
         data = {
             "userId": user_id,
-            "deviceName": device_name,
             "skew": skew if skew is not None else self.config.default_skew,
             "period": period if period is not None else self.config.default_period,
         }
+        if device_name is not None:
+            data["deviceName"] = device_name
         response = await self.querier.send_post_request(
             NormalisedURLPath("/recipe/totp/device"),
             data,
diff --git a/tests/auth-react/flask-server/app.py b/tests/auth-react/flask-server/app.py
index ddb3b4b7a..913183655 100644
--- a/tests/auth-react/flask-server/app.py
+++ b/tests/auth-react/flask-server/app.py
@@ -174,7 +174,7 @@ def get_website_domain():
 
 os.environ.setdefault("SUPERTOKENS_ENV", "testing")
 
-latest_url_with_token = None
+latest_url_with_token = ""
 
 code_store: Dict[str, List[Dict[str, Any]]] = {}
 accountlinking_config: Dict[str, Any] = {}
@@ -1286,6 +1286,8 @@ def before_each():
     global enabled_providers
     global enabled_recipes
     global mfa_info
+    global latest_url_with_token
+    latest_url_with_token = ""
     code_store = dict()
     accountlinking_config = {}
     enabled_providers = None
@@ -1319,6 +1321,16 @@ def test_set_account_linking_config():
     return "", 200
 
 
+@app.route("/setMFAInfo", methods=["POST"])  # type: ignore
+def set_mfa_info():
+    global mfa_info
+    body = request.get_json()
+    if body is None:
+        return jsonify({"error": "Invalid request body"}), 400
+    mfa_info = body
+    return jsonify({"status": "OK"})
+
+
 @app.route("/test/setEnabledRecipes", methods=["POST"])  # type: ignore
 def test_set_enabled_recipes():
     global enabled_recipes
@@ -1332,6 +1344,21 @@ def test_set_enabled_recipes():
     return "", 200
 
 
+@app.route("/test/getTOTPCode", methods=["POST"])  # type: ignore
+def test_get_totp_code():
+    from pyotp import TOTP
+
+    body = request.get_json()
+    if body is None or "secret" not in body:
+        return jsonify({"error": "Invalid request body"}), 400
+
+    secret = body["secret"]
+    totp = TOTP(secret, digits=6, interval=1)
+    code = totp.now()
+
+    return jsonify({"totp": code})
+
+
 @app.get("/test/getDevice")  # type: ignore
 def test_get_device():
     global code_store