Add Models to JWTRoutes class & init_app method #119

Author: joegasewicz

Pipfile

@@ -8,6 +8,7 @@ pytest = "*"
 twine = "*"
 pyjwt = "*"
 Flask-SQLAlchemy = "*"
+Sphinx = "*"
 
 [packages]
 flask = "*"

flask_jwt_router/_authentication.py

@@ -38,6 +38,7 @@ def update_entity(self, extensions: Config, exp: int, table_name, **kwarg):
         # pylint:disable=missing-function-docstring
         pass
 
+    @abstractmethod
     def encode_token(self, extensions: Config, entity_id: Any, exp: int, table_name: str):
         # pylint:disable=missing-function-docstring
         pass

flask_jwt_router/_extensions.py

@@ -2,7 +2,9 @@
      The main configuration class for Flask-JWT-Router
 """
 from abc import ABC
-from typing import Dict, Any
+from typing import Dict, Any, List
+
+from ._entity import _ORMType
 
 
 class Config:
@@ -33,23 +35,26 @@ def __init__(self,
 
 class BaseExtension(ABC):
     """Abstract Base Class for Extensions"""
-    def init_extensions(self, config: Dict[str, Any]) -> Config:
+    def init_extensions(self, config: Dict[str, Any], **kwargs) -> Config:
         # pylint: disable=missing-function-docstring
         pass
 
 
 class Extensions(BaseExtension):
     """Contains the main configuration values"""
-    def init_extensions(self, config: Any) -> Config:
+    entity_models: List[_ORMType]
+
+    def init_extensions(self, config: Any, **kwargs) -> Config:
         """
         :param config:
         :return:
         """
+        entity_models = kwargs.get("entity_models")
         return Config(
             config.get("SECRET_KEY") or "DEFAULT_SECRET_KEY",
             config.get("ENTITY_KEY") or "id",
             config.get("WHITE_LIST_ROUTES") or [],
             config.get("JWT_ROUTER_API_NAME"),
             config.get("IGNORED_ROUTES") or [],
-            config.get("ENTITY_MODELS") or [],
+            entity_models or config.get("ENTITY_MODELS") or [],
         )

flask_jwt_router/_jwt_router.py

@@ -4,9 +4,10 @@
 
 import logging
 from warnings import warn
+from typing import List
 
 from ._extensions import BaseExtension, Extensions, Config
-from ._entity import BaseEntity, Entity
+from ._entity import BaseEntity, Entity, _ORMType
 from ._routing import BaseRouting, Routing
 from ._authentication import BaseAuthStrategy
 
@@ -27,6 +28,9 @@ class FlaskJWTRouter:
     #: The Flask application instance.
     app = None
 
+    #: A list of entity models
+    entity_models: List[_ORMType]
+
     #: Token expiry value. eg. 30 = 30 days from creation date.
     exp: int = 30
 
@@ -50,20 +54,23 @@ class FlaskJWTRouter:
     #: for more information.
     ext: BaseExtension
 
-    def __init__(self, app=None):
+    def __init__(self, app=None, **kwargs):
+        self.entity_models = kwargs.get("entity_models")
         self.ext = Extensions()
+        self.app = app
         if app:
-            self.init_app(app)
+            self.init_app(app, entity_models=self.entity_models)
 
-    def init_app(self, app):
+    def init_app(self, app=None, **kwargs):
         """
         You can use this to set up your config at runtime
         :param app: Flask application instance
         :return:
         """
-        self.app = app
+        self.app = self.app or app
+        entity_models = self.entity_models or kwargs.get("entity_models")
         config = self.get_app_config(app)
-        self.extensions = self.ext.init_extensions(config)
+        self.extensions = self.ext.init_extensions(config, entity_models=entity_models)
         self.entity = Entity(self.extensions)
         self.routing = Routing(self.app, self.extensions, self.entity)
         self.app.before_request(self.routing.before_middleware)

tests/test_jwt_routes.py

@@ -0,0 +1,36 @@
+from flask import Flask
+
+
+from flask_jwt_router._jwt_routes import JwtRoutes
+from tests.fixtures.model_fixtures import MockEntityModel
+
+
+class TestJwtRoutes:
+
+    app = Flask(__name__)
+
+    def test_init_app(self):
+        jwt = JwtRoutes(self.app, entity_models=[MockEntityModel])
+        assert jwt.extensions.entity_models[0] == MockEntityModel
+        assert jwt.app == self.app
+
+        jwt = JwtRoutes()
+        jwt.init_app(self.app, entity_models=[MockEntityModel])
+        assert jwt.extensions.entity_models[0] == MockEntityModel
+        assert jwt.app == self.app
+
+        jwt = JwtRoutes(self.app)
+        jwt.init_app(entity_models=[MockEntityModel])
+        assert jwt.extensions.entity_models[0] == MockEntityModel
+        assert jwt.app == self.app
+
+        jwt = JwtRoutes(entity_models=[MockEntityModel])
+        jwt.init_app(self.app)
+        assert jwt.extensions.entity_models[0] == MockEntityModel
+        assert jwt.app == self.app
+
+        self.app.config["ENTITY_MODELS"] = [MockEntityModel]
+        jwt = JwtRoutes()
+        jwt.init_app(self.app)
+        assert jwt.extensions.entity_models[0] == MockEntityModel
+        assert jwt.app == self.app