Add JWT Authentication to API Endpoints

ComPlat · Sep 29, 2023 · 19d0816 · 19d0816
1 parent 5b5f57e
commit 19d0816
Show file tree

Hide file tree

Showing 8 changed files with 94 additions and 49 deletions.
diff --git a/chem_spectra/__init__.py b/chem_spectra/__init__.py
@@ -1,6 +1,7 @@
 import os
 
 from flask import Flask
+from flask_jwt_extended import JWTManager
 
 import logging
 
@@ -23,6 +24,8 @@ def create_app(test_config=None):
     except OSError:
         pass
 
+    jwt = JWTManager(app)
+
     #create logging
     logger = logging.getLogger(__name__)
     logger.setLevel(logging.INFO)
@@ -56,6 +59,7 @@ def ping():
     from chem_spectra.controller.spectra_type_api import spectra_type_api
     app.register_blueprint(spectra_type_api)
 
+    # spectra layout api
     from chem_spectra.controller.spectra_layout_api import spectra_layout_api
     app.register_blueprint(spectra_layout_api)
 

diff --git a/chem_spectra/controller/spectra_layout_api.py b/chem_spectra/controller/spectra_layout_api.py
@@ -2,6 +2,7 @@
 from flask import jsonify, Blueprint
 import os
 import shutil
+from flask_jwt_extended import jwt_required
 
 spectra_layout_api = Blueprint('spectra_layout_api', __name__)
 script_dir = os.path.dirname(__file__)
@@ -17,7 +18,8 @@ def load_data_types():
         with open(data_type_json_path, 'r') as mapping_file:
             return json.load(mapping_file)
 
-@spectra_layout_api.route('/spectra_layouts', methods=['GET'])
+@spectra_layout_api.route('/api/v1/chemspectra/spectra_layouts', methods=['GET'])
+@jwt_required()
 def get_spectra_layouts():
     existing_data_types = load_data_types()
     return jsonify(existing_data_types["datatypes"]), 200
diff --git a/chem_spectra/controller/spectra_type_api.py b/chem_spectra/controller/spectra_type_api.py
@@ -2,6 +2,7 @@
 import os
 from flask import request, jsonify, Blueprint
 import shutil
+from flask_jwt_extended import jwt_required
 
 spectra_type_api = Blueprint('spectra_type_api', __name__)
 
@@ -22,7 +23,8 @@ def save_data_types(data_types):
     with open(data_type_json_path, 'w') as mapping_file:
         json.dump(data_types, mapping_file, indent=4)
 
-@spectra_type_api.route('/data_type', methods=['POST'])
+@spectra_type_api.route('/api/v1/chemspectra/data_type', methods=['POST'])
+@jwt_required()
 def create_data_type():
     request_data = request.get_json()
     new_data_type_mapping = request_data.get("new_data_type")
@@ -40,7 +42,8 @@ def create_data_type():
 
     return jsonify({"message": "Data type created successfully"}), 200
 
-@spectra_type_api.route('/data_type/<data_type>', methods=['DELETE'])
+@spectra_type_api.route('/api/v1/chemspectra/data_type/<data_type>', methods=['DELETE'])
+@jwt_required()
 def delete_data_type(data_type):
     existing_data_types = load_data_types()
 

diff --git a/requirements.txt b/requirements.txt
@@ -42,3 +42,5 @@ urllib3==1.26.5
 Werkzeug==2.2.3
 zipp==0.5.2
 pyopenms==2.6.0
+PyJWT==2.8.0
+flask-jwt-extended==4.5.2
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -6,7 +6,8 @@
 def app():
     app = create_app({
         'TESTING': True,
-        'IP_WHITE_LIST': '127.0.0.1'
+        'IP_WHITE_LIST': '127.0.0.1',
+        'JWT_SECRET_KEY': 'your_secret_key'
     })
 
     yield app

diff --git a/tests/controller/test_spectra_layout_api.py b/tests/controller/test_spectra_layout_api.py
@@ -1,14 +1,21 @@
 import os
 import json
+from unittest.mock import patch
+from instance.config import TEST_TOKEN
+
+test_json_path = './tests/fixtures/test_data_types.json'
+orig_json_path = 'chem_spectra.controller.spectra_layout_api.data_type_json_path'
+test_jwt_token = TEST_TOKEN
 
 def test_get_spectra_layouts_with_data(client):
-    data_type_json_path = './tests/fixtures/test_data_types.json'
-    response = client.get('/spectra_layouts')
-    response_data = response.json()
-    assert response.status_code == 200
-    assert response_data == {
-        "INFRARED": ["INFRARED SPECTRUM"],
-        "MS": ["MASS SPECTRUM"],
-        "NMR": ["NMR SPECTRUM", "NMRSPECTRUM"],
-        "RAMAN": ["RAMAN"]
-        }
+    with patch(orig_json_path, new=test_json_path):
+        response = client.get('/api/v1/chemspectra/spectra_layouts', 
+                              headers={'Authorization': f'Bearer {test_jwt_token}'})
+        response_data = response.json
+        assert response.status_code == 200
+        assert response_data == {
+            "INFRARED": ["INFRARED SPECTRUM"],
+            "MS": ["MASS SPECTRUM"],
+            "NMR": ["NMR SPECTRUM", "NMRSPECTRUM"],
+            "RAMAN": ["RAMAN"]
+        }
diff --git a/tests/controller/test_spectra_type_api.py b/tests/controller/test_spectra_type_api.py
@@ -1,21 +1,31 @@
 import os
 import json
+import unittest
+from unittest.mock import patch
+from instance.config import TEST_TOKEN
 
-data_type_json_path = './tests/fixtures/test_data_types.json'
+test_json_path = './tests/fixtures/test_data_types.json'
+orig_json_path = 'chem_spectra.controller.spectra_type_api.data_type_json_path'
+test_jwt_token = TEST_TOKEN
 
 def test_create_data_type(client):
     new_data_type = {
         "new_data_type": {
             "MS": "MASS SPEC"
         }
     }
-
-    response = client.post('/data_type', json=new_data_type)
-    assert response.status_code == 200
-
-    response_data = response.json()
-    assert "message" in response_data
-    response_data["message"] == "Data type created successfully"
+
+    with patch(orig_json_path, test_json_path):
+        response = client.post('/api/v1/chemspectra/data_type', json=new_data_type, 
+                               headers={'Authorization': f'Bearer {test_jwt_token}'})
+
+        assert response.status_code == 200
+
+        response_data = response.json
+        assert "message" in response_data
+        response_data["message"] == "Data type created successfully"
+        response = client.delete('/api/v1/chemspectra/data_type/MASS SPEC', 
+                                 headers={'Authorization': f'Bearer {test_jwt_token}'})
 
 def test_create_data_type_unchanged(client):    
     # data type already exists in JSON
@@ -25,13 +35,15 @@ def test_create_data_type_unchanged(client):
         }
     }
 
-    response = client.post('/data_type', json=new_data_type)
-    assert response.status_code == 400
+    with patch(orig_json_path, test_json_path):
+        response = client.post('/api/v1/chemspectra/data_type', json=new_data_type, 
+                               headers={'Authorization': f'Bearer {test_jwt_token}'})
+        assert response.status_code == 400
 
-    response_data = response.json()
-
-    assert "message" in response_data
-    assert response_data["message"] == "Data type 'INFRARED SPECTRUM' already exists"
+        response_data = response.json
+        
+        assert "message" in response_data
+        assert response_data["message"] == "Data type 'INFRARED SPECTRUM' already exists"
 
 def test_create_data_type_layout_does_not_exist(client):
     new_data_type = {
@@ -40,13 +52,15 @@ def test_create_data_type_layout_does_not_exist(client):
         }
     }
 
-    response = client.post('/data_type', json=new_data_type)
-    assert response.status_code == 400
+    with patch(orig_json_path, test_json_path):
+        response = client.post('/api/v1/chemspectra/data_type', json=new_data_type, 
+                               headers={'Authorization': f'Bearer {test_jwt_token}'})
+        assert response.status_code == 400
 
-    response_data = response.json()
-
-    assert "message" in response_data
-    assert response_data["message"] == "Layout 'UNKNOWN' does not exist"
+        response_data = response.json
+        
+        assert "message" in response_data
+        assert response_data["message"] == "Layout 'UNKNOWN' does not exist"
 
 def test_delete_data_type(client):
     # create a new data type
@@ -55,13 +69,16 @@ def test_delete_data_type(client):
             "INFRARED": "IR"
         }
     }
-    response = client.post('/data_type', json=new_data_type)
-    assert response.status_code == 200
+    with patch(orig_json_path, test_json_path):
+        response = client.post('/api/v1/chemspectra/data_type', json=new_data_type, 
+                               headers={'Authorization': f'Bearer {test_jwt_token}'})
+        assert response.status_code == 200
 
-    # delete the new data type
-    response = client.delete('/data_type/IR')
-    assert response.status_code == 200
-    response_data = response.json()
+        # delete the new data type
+        response = client.delete('/api/v1/chemspectra/data_type/IR', 
+                                 headers={'Authorization': f'Bearer {test_jwt_token}'})
+        assert response.status_code == 200
+        response_data = response.json
 
-    assert "message" in response_data
-    assert response_data["message"] == "Data type 'IR' deleted successfully"
+        assert "message" in response_data
+        assert response_data["message"] == "Data type 'IR' deleted successfully"
diff --git a/tests/fixtures/test_data_types.json b/tests/fixtures/test_data_types.json
@@ -1,8 +1,17 @@
 {
-  "datatypes": {
-    "NMR": ["NMR SPECTRUM", "NMRSPECTRUM"],
-    "INFRARED": ["INFRARED SPECTRUM"],
-    "RAMAN": ["RAMAN"],
-    "MS": ["MASS SPECTRUM"]
-  }
-}
+    "datatypes": {
+        "NMR": [
+            "NMR SPECTRUM",
+            "NMRSPECTRUM"
+        ],
+        "INFRARED": [
+            "INFRARED SPECTRUM"
+        ],
+        "RAMAN": [
+            "RAMAN"
+        ],
+        "MS": [
+            "MASS SPECTRUM"
+        ]
+    }
+}