test: Fix issues with tests

Author: KShivendu

tests/Django/test_django.py

@@ -13,23 +13,24 @@
 # under the License.
 
 import json
+from datetime import datetime
 from inspect import isawaitable
 from typing import Any, Dict, Union
 
-from datetime import datetime
 from django.http import HttpRequest, HttpResponse, JsonResponse
 from django.test import RequestFactory, TestCase
+from pytest import mark
 from supertokens_python import InputAppInfo, SupertokensConfig, init
 from supertokens_python.framework.django import middleware
 from supertokens_python.recipe import emailpassword, session
-from supertokens_python.recipe.emailpassword.interfaces import APIInterface, APIOptions
+from supertokens_python.recipe.emailpassword.interfaces import (APIInterface,
+                                                                APIOptions)
 from supertokens_python.recipe.session import SessionContainer
-from supertokens_python.recipe.session.asyncio import (
-    create_new_session,
-    get_session,
-    refresh_session,
-)
-from supertokens_python.recipe.session.framework.django.asyncio import verify_session
+from supertokens_python.recipe.session.asyncio import (create_new_session,
+                                                       get_session,
+                                                       refresh_session)
+from supertokens_python.recipe.session.framework.django.asyncio import \
+    verify_session
 from tests.utils import clean_st, reset, setup_st, start_st
 
 
@@ -389,3 +390,19 @@ async def test_optional_session(self):
         assert response.status_code == 200
         dict_response = json.loads(response.content)
         assert dict_response["s"] == "empty session"
+
+
+from django.http import HttpResponse
+from supertokens_python.framework.django.django_response import \
+    DjangoResponse as SuperTokensDjangoWrapper
+
+
+class SupertokensResponseTest(TestCase):
+    def test_remove_header_works(self):
+        response = HttpResponse()
+        st_response = SuperTokensDjangoWrapper(response)
+
+        st_response.set_header("foo", "bar")
+        assert st_response.get_header("foo") == "bar"
+        st_response.remove_header("foo")
+        assert st_response.get_header("foo") is None

tests/Flask/test_flask.py

@@ -398,3 +398,15 @@ def test_optional_session(driver_config_app: Any):
     dict_response = json.loads(response.data)
     assert response.status_code == 200
     assert dict_response["s"] == "empty session"
+
+from flask import Response
+from supertokens_python.framework.flask.flask_response import FlaskResponse as SupertokensFlaskWrapper
+
+def test_remove_header_works():
+    response = Response()
+    st_response = SupertokensFlaskWrapper(response)
+
+    st_response.set_header("foo", "bar")
+    assert st_response.get_header("foo") == "bar"
+    st_response.remove_header("foo")
+    assert st_response.get_header("foo") is None

tests/utils.py

@@ -19,7 +19,7 @@
 from signal import SIGTERM
 from subprocess import DEVNULL, run
 from time import sleep
-from typing import Any, Dict, List, cast
+from typing import Any, Dict, List, cast, Optional
 
 from requests.models import Response
 from yaml import FullLoader, dump, load
@@ -252,8 +252,8 @@ def extract_info(response: Response) -> Dict[str, Any]:
         "antiCsrf": response.headers.get("anti-csrf"),
         "accessTokenFromHeader": access_token_from_header,
         "refreshTokenFromHeader": refresh_token_from_header,
-        "accessTokenFromAny": access_token_from_header or access_token,
-        "refreshTokenFromAny": refresh_token_from_header or refresh_token,
+        "accessTokenFromAny": access_token_from_header if access_token is None else access_token,
+        "refreshTokenFromAny": refresh_token_from_header if refresh_token is None else refresh_token,
     }
 
 
@@ -517,10 +517,14 @@ async def __call__(  # pylint: disable=invalid-overridden-method, useless-super-
 }
 
 
-def get_st_init_args(recipe_list: List[Any]) -> Dict[str, Any]:
+def get_st_init_args(recipe_list: List[Any], extra_config: Optional[Dict[str, Any]]=None) -> Dict[str, Any]:
+    if extra_config is None:
+        extra_config = {}
+
     return {
         **st_init_common_args,
         "recipe_list": recipe_list,
+        **extra_config
     }