From 1dd21a3882508e03fdf45b38e5bdfa28e8edd6c4 Mon Sep 17 00:00:00 2001
From: Tochukwu <dlastdon4000@gmail.com>
Date: Sat, 27 Apr 2024 06:56:56 +0100
Subject: [PATCH] Fix: Verify and Blacklist custom response schema (#72)

* fixed verify and blacklist response schema #71

* 5.3.1
---
 ninja_jwt/__init__.py          |  2 +-
 ninja_jwt/controller.py        | 12 ++++++------
 ninja_jwt/routers/blacklist.py |  3 +--
 ninja_jwt/routers/verify.py    |  3 +--
 tests/test_custom_schema.py    | 23 +++++++++++++++++++++--
 5 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/ninja_jwt/__init__.py b/ninja_jwt/__init__.py
index 5892b6b81..21dd23349 100644
--- a/ninja_jwt/__init__.py
+++ b/ninja_jwt/__init__.py
@@ -1,3 +1,3 @@
 """Django Ninja JWT - JSON Web Token for Django-Ninja"""
 
-__version__ = "5.3.0"
+__version__ = "5.3.1"
diff --git a/ninja_jwt/controller.py b/ninja_jwt/controller.py
index 863c49d61..996f036c5 100644
--- a/ninja_jwt/controller.py
+++ b/ninja_jwt/controller.py
@@ -1,5 +1,4 @@
 from asgiref.sync import sync_to_async
-from ninja import Schema
 from ninja_extra import ControllerBase, api_controller, http_post
 from ninja_extra.permissions import AllowAny
 
@@ -31,11 +30,12 @@ class TokenVerificationController(ControllerBase):
 
     @http_post(
         "/verify",
-        response={200: Schema},
+        response={200: schema.verify_schema.get_response_schema()},
         url_name="token_verify",
     )
     def verify_token(self, token: schema.verify_schema):
-        return token.to_response_schema()
+        asas = token.to_response_schema()
+        return asas
 
 
 class TokenBlackListController(ControllerBase):
@@ -43,7 +43,7 @@ class TokenBlackListController(ControllerBase):
 
     @http_post(
         "/blacklist",
-        response={200: Schema},
+        response={200: schema.blacklist_schema.get_response_schema()},
         url_name="token_blacklist",
     )
     def blacklist_token(self, refresh: schema.blacklist_schema):
@@ -114,7 +114,7 @@ class NinjaJWTSlidingController(
 class AsyncTokenVerificationController(TokenVerificationController):
     @http_post(
         "/verify",
-        response={200: Schema},
+        response={200: schema.verify_schema.get_response_schema()},
         url_name="token_verify",
     )
     async def verify_token(self, token: schema.verify_schema):
@@ -126,7 +126,7 @@ class AsyncTokenBlackListController(TokenBlackListController):
 
     @http_post(
         "/blacklist",
-        response={200: Schema},
+        response={200: schema.blacklist_schema.get_response_schema()},
         url_name="token_blacklist",
     )
     async def blacklist_token(self, refresh: schema.blacklist_schema):
diff --git a/ninja_jwt/routers/blacklist.py b/ninja_jwt/routers/blacklist.py
index 450bf27e6..d7644db4d 100644
--- a/ninja_jwt/routers/blacklist.py
+++ b/ninja_jwt/routers/blacklist.py
@@ -1,4 +1,3 @@
-from ninja import Schema
 from ninja.router import Router
 
 from ninja_jwt.schema_control import SchemaControl
@@ -11,7 +10,7 @@
 
 @blacklist_router.post(
     "/blacklist",
-    response={200: Schema},
+    response={200: schema.blacklist_schema.get_response_schema()},
     url_name="token_blacklist",
     auth=None,
 )
diff --git a/ninja_jwt/routers/verify.py b/ninja_jwt/routers/verify.py
index 8b851f8da..8a684dc33 100644
--- a/ninja_jwt/routers/verify.py
+++ b/ninja_jwt/routers/verify.py
@@ -1,4 +1,3 @@
-from ninja import Schema
 from ninja.router import Router
 
 from ninja_jwt.schema_control import SchemaControl
@@ -11,7 +10,7 @@
 
 @verify_router.post(
     "/verify",
-    response={200: Schema},
+    response={200: schema.verify_schema.get_response_schema()},
     url_name="token_verify",
     auth=None,
 )
diff --git a/tests/test_custom_schema.py b/tests/test_custom_schema.py
index 554adc0ad..3dc70b14b 100644
--- a/tests/test_custom_schema.py
+++ b/tests/test_custom_schema.py
@@ -91,7 +91,22 @@ def get_response_schema(cls):
 
 
 class MyTokenVerifyInputSchema(TokenVerifyInputSchema):
-    pass
+    @classmethod
+    def get_response_schema(cls):
+        class NewResponseSchema(Schema):
+            refresh: str
+            access: str
+            user: dict
+
+        return NewResponseSchema
+
+    def to_response_schema(self):
+        values = {
+            "refresh": "your_refresh_token_here",
+            "access": self.token,
+            "user": {},
+        }
+        return values
 
 
 class MyTokenBlacklistInputSchema(TokenBlacklistInputSchema):
@@ -330,7 +345,11 @@ def test_it_should_return_200_if_everything_okay(self, monkeypatch):
                 "/verify", json={"token": str(token)}, content_type="application/json"
             )
         assert res.status_code == 200
-        assert res.json() == {}
+        data = res.json()
+
+        assert "refresh" in data
+        assert "user" in data
+        assert "access" in data
 
 
 @pytest.mark.django_db