Add Models to JWTRoutes class & init_app method #119

Pipfile

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

README.md

@@ -78,14 +78,17 @@ def register():
 class UserModel(db.Model):
     id = db.Column(db.Integer, primary_key=True)
     name = db.Column(db.String)
+
 # You can define the primary key name with `ENTITY_KEY` on Flask's config
 app.config["ENTITY_KEY"] = "user_id"
-# (`id` is used by default)
-JwtRoutes(app)
 
-# You can also specify a list of entity model classes
+# (`id` is used by default)
+JwtRoutes(app, entity_models=[UserModel, TeacherModel, ...etc])
 
-app.config["ENTITY_MODELS"] = [ UserModel, TeacherModel, ...etc ]
+# Or pass later with `init_app`
+def create_app(config):
+    ...
+    jwt_routes.init_app(app, entity_models=[UserModel, TeacherModel, ...etc])
 
 ```
 
@@ -181,7 +184,9 @@ An Example configuration for registering & logging in users of different types:
         ("POST", "/auth/user"), ("POST", "/auth/user/login"),
         ("POST", "/auth/teacher"), ("POST", "/auth/teacher/login"),
     ]
-    app.config["ENTITY_MODELS"] = [UserModel, TeacherModel]
+
+    # Optionally, you can pass your models to Flask's config:
+    app.config["ENTITY_MODELS"] = [ UserModel, TeacherModel, ...etc ]
 ```
 ## Authors
 

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

@@ -1,8 +1,10 @@
 """
      The main configuration class for Flask-JWT-Router
 """
-from abc import ABC
-from typing import Dict, Any
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List
+
+from ._entity import _ORMType
 
 
 class Config:
@@ -33,23 +35,27 @@ def __init__(self,
 
 class BaseExtension(ABC):
     """Abstract Base Class for Extensions"""
-    def init_extensions(self, config: Dict[str, Any]) -> Config:
+    @abstractmethod
+    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)

flask_jwt_router/_jwt_routes.py

@@ -66,15 +66,16 @@ def register():
         # Create your entity model (example uses Flask-SqlAlchemy)
 
         class UserModel(db.Model):
+            __tablename__ = "users"
             id = db.Column(db.Integer, primary_key=True)
             name = db.Column(db.String)
 
-        # You can also specify a list of entity model classes
+        JwtRoutes(app, entity_models=[UserModel, TeacherModel, ...etc])
 
-        app.config["ENTITY_MODELS"] = [ UserModel, TeacherModel ]
-
-        # (`id` is used by default)
-        JwtRoutes(app)
+        # Or pass later with `init_app`
+        def create_app(config):
+            ...
+            jwt_routes.init_app(app, entity_models=[UserModel, TeacherModel, ...etc])
 
 
     Authorization & Tokens

tests/test_extensions.py

@@ -1,29 +1,32 @@
 from flask_jwt_router._jwt_routes import JwtRoutes
+from flask_jwt_router._extensions import Extensions
+from tests.fixtures.model_fixtures import MockEntityModel
 
 
 class TestExtension:
+    IGNORED_ROUTES = [
+        ("GET", "/"),
+        ("GET", "/ignore"),
+    ]
+    WHITE_LIST_ROUTES = [
+        ("GET", "/test"),
+    ]
 
-    def test_config(self):
-        IGNORED_ROUTES = [
-            ("GET", "/"),
-            ("GET", "/ignore"),
-        ]
-        WHITE_LIST_ROUTES = [
-            ("GET", "/test"),
-        ]
-        class App:
-            config = {
+    config = {
                 "IGNORED_ROUTES": IGNORED_ROUTES,
                 "WHITE_LIST_ROUTES": WHITE_LIST_ROUTES,
             }
-            def before_request(self, t):
-                pass
-        flask_jwt_router = JwtRoutes(App())
-        assert flask_jwt_router.extensions.ignored_routes == IGNORED_ROUTES
-        assert flask_jwt_router.extensions.whitelist_routes == WHITE_LIST_ROUTES
-        flask_jwt = JwtRoutes()
-        flask_jwt.init_app(App())
-        assert flask_jwt.extensions.ignored_routes == IGNORED_ROUTES
-        assert flask_jwt.extensions.whitelist_routes == WHITE_LIST_ROUTES
 
+    def test_init_extensions(self, MockEntityModel):
+        extensions = Extensions()
+        config = extensions.init_extensions(self.config, entity_models=[MockEntityModel])
+
+        assert config.whitelist_routes == self.WHITE_LIST_ROUTES
+        assert config.ignored_routes == self.IGNORED_ROUTES
+        assert config.entity_models == [MockEntityModel]
+
+        config = {**self.config, "ENTITY_MODELS": [MockEntityModel]}
+        con = extensions.init_extensions(config)
+
+        assert con.entity_models == [MockEntityModel]
 

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, MockEntityModel):
+        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