build(tests): lint, unit tests, and ci for that (data-platform-hq#32)

Co-authored-by: Aleksandr Chikovani <[email protected]>
Co-authored-by: Alexander Kharkevich <[email protected]>

.coveragerc

@@ -1,3 +1,3 @@
 [run]
 source = ./
-omit = mlflow_oidc_auth/tests/*,mlflow_oidc_auth/db/migrations/versions/*
+omit = mlflow_oidc_auth/tests/*,mlflow_oidc_auth/db/migrations/versions/*,mlflow_oidc_auth/views/*

.github/workflows/pre-commit.yml

@@ -0,0 +1,27 @@
+name: Run pre-commit
+on:
+  pull_request:
+    types:
+      - opened
+      - edited
+      - reopened
+      - synchronize
+jobs:
+  pre-commit:
+    name: Run pre-commit
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        with:
+          fetch-depth: 0
+      - name: Set up Python
+        uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5.1.1
+        with:
+          python-version: 3.11
+      - name: Run pre-commit
+        run: |
+          python -m pip install --upgrade pip
+          pip install pre-commit
+          pre-commit install
+          pre-commit run --all-files --show-diff-on-failure

.github/workflows/unit-tests.yml

@@ -0,0 +1,35 @@
+name: Unit tests
+on:
+  pull_request:
+    types:
+      - opened
+      - edited
+      - reopened
+      - synchronize
+  push:
+    branches:
+      - main
+jobs:
+  python-tests:
+    name: Run python tests
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332 # v4.1.7
+        with:
+          fetch-depth: 0
+      - name: Set up Python
+        uses: actions/setup-python@39cd14951b08e74b54015e9e001cdefcf80e669f # v5.1.1
+        with:
+          python-version: 3.11
+      - name: Run tests
+        run: |
+          python -m pip install --upgrade pip
+          pip install tox
+          tox -e py
+      - name: Override Coverage Source Path for Sonar
+        run: sed -i "s@<source>/home/runner/work/mlflow-oidc-auth/mlflow-oidc-auth</source>@<source>/github/workspace</source>@g" /home/runner/work/mlflow-oidc-auth/mlflow-oidc-auth/coverage.xml
+      - name: SonarCloud Scan
+        uses: SonarSource/sonarcloud-github-action@e44258b109568baa0df60ed515909fc6c72cba92 # v2.3.0
+        env:
+          SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }}

.gitignore

@@ -176,3 +176,4 @@ flask_session/
 node_modules/
 .angular/
 mlflow_oidc_auth/ui
+pytest-coverage.txt

README.md

@@ -33,7 +33,6 @@ The plugin required the following environment variables but also supported `.env
 | OIDC_TOKEN_URL         | OIDC Token URL (if discovery URL is not defined) |
 | OIDC_USER_URL          | OIDC User info URL (if discovery URL is not defined) |
 | SECRET_KEY             | Key to perform cookie encryption |
-| OAUTHLIB_INSECURE_TRANSPORT | Development only. Allow to use insecure endpoints for OIDC |
 | LOG_LEVEL                   | Application log level |
 | OIDC_USERS_DB_URI | Database connection string |
 

docs/configuration/examples/microsoft.md

@@ -19,4 +19,3 @@ OIDC_SCOPE = "openid,profile,email"
 OIDC_GROUP_NAME = "mlflow_users_group_name"
 OIDC_ADMIN_GROUP_NAME = "mlflow_admins_group_name"
 ```
-

docs/configuration/index.md

@@ -17,7 +17,6 @@ The plugin required the following environment variables but also supported `.env
 | OIDC_TOKEN_URL         | OIDC Token URL (if discovery URL is not defined) |
 | OIDC_USER_URL          | OIDC User info URL (if discovery URL is not defined) |
 | SECRET_KEY             | Key to perform cookie encryption |
-| OAUTHLIB_INSECURE_TRANSPORT | Development only. Allow to use insecure endpoints for OIDC |
 | LOG_LEVEL                   | Application log level |
 | OIDC_USERS_DB_URI | Database connection string |
 
@@ -35,5 +34,3 @@ The plugin required the following environment variables but also supported `.env
 | REDIS_USERNAME | Redis username | None |
 | REDIS_PASSWORD | Redis password | None |
 | REDIS_SSL | Use SSL | false |
-
-

docs/permission-management/index.md

@@ -4,4 +4,3 @@
 
 
 ## Permissions hierarchy
-

mlflow_oidc_auth/auth.py

@@ -1,26 +1,38 @@
-from typing import Union
+from typing import Union, Optional
 
 import requests
 from authlib.integrations.flask_client import OAuth
 from authlib.jose import jwt
 from flask import Response, request
 from werkzeug.datastructures import Authorization
 
-from mlflow_oidc_auth.app import app
 from mlflow_oidc_auth.config import config
 from mlflow_oidc_auth.store import store
-oauth = OAuth(app)
 
-oauth.register(
-    name="oidc",
-    client_id=config.OIDC_CLIENT_ID,
-    client_secret=config.OIDC_CLIENT_SECRET,
-    server_metadata_url=config.OIDC_DISCOVERY_URL,
-    client_kwargs={"scope": config.OIDC_SCOPE},
-)
+
+_oauth_instance: Optional[OAuth] = None
+
+
+def get_oauth_instance(app) -> OAuth:
+    # returns a singleton instance of OAuth
+    # to avoid circular imports
+    global _oauth_instance
+
+    if _oauth_instance is None:
+        _oauth_instance = OAuth(app)
+        _oauth_instance.register(
+            name="oidc",
+            client_id=config.OIDC_CLIENT_ID,
+            client_secret=config.OIDC_CLIENT_SECRET,
+            server_metadata_url=config.OIDC_DISCOVERY_URL,
+            client_kwargs={"scope": config.OIDC_SCOPE},
+        )
+    return _oauth_instance
+
 
 def _get_oidc_jwks():
-    from mlflow_oidc_auth.app import cache
+    from mlflow_oidc_auth.app import cache, app
+
     jwks = cache.get("jwks")
     if jwks:
         app.logger.debug("JWKS cache hit")
@@ -41,6 +53,8 @@ def validate_token(token):
 
 
 def authenticate_request_basic_auth() -> Union[Authorization, Response]:
+    from mlflow_oidc_auth.app import app
+
     username = request.authorization.username
     password = request.authorization.password
     app.logger.debug("Authenticating user %s", username)
@@ -53,6 +67,8 @@ def authenticate_request_basic_auth() -> Union[Authorization, Response]:
 
 
 def authenticate_request_bearer_token() -> Union[Authorization, Response]:
+    from mlflow_oidc_auth.app import app
+
     token = request.authorization.token
     try:
         user = validate_token(token)

mlflow_oidc_auth/config.py

@@ -10,6 +10,7 @@
 load_dotenv()  # take environment variables from .env.
 app.logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))
 
+
 class AppConfig:
     def __init__(self):
         self.DEFAULT_MLFLOW_PERMISSION = os.environ.get("DEFAULT_MLFLOW_PERMISSION", "MANAGE")
@@ -52,4 +53,5 @@ def __init__(self):
             except ImportError:
                 app.logger.error(f"Cache module for {self.CACHE_TYPE} could not be imported.")
 
+
 config = AppConfig()

mlflow_oidc_auth/sqlalchemy_store.py

@@ -113,6 +113,7 @@ def delete_user(self, username: str):
         with self.ManagedSessionMaker() as session:
             user = self._get_user(session, username)
             session.delete(user)
+            session.flush()
 
     def create_experiment_permission(self, experiment_id: str, username: str, permission: str) -> ExperimentPermission:
         _validate_permission(permission)
@@ -341,6 +342,7 @@ def delete_registered_model_permission(self, name: str, username: str):
         with self.ManagedSessionMaker() as session:
             perm = self._get_registered_model_permission(session, name, username)
             session.delete(perm)
+            session.flush()
 
     def list_experiment_permissions_for_experiment(self, experiment_id: str):
         with self.ManagedSessionMaker() as session:

mlflow_oidc_auth/tests/test_auth.py

@@ -0,0 +1,111 @@
+from unittest.mock import patch, MagicMock
+from mlflow_oidc_auth.auth import (
+    get_oauth_instance,
+    _get_oidc_jwks,
+    validate_token,
+    authenticate_request_basic_auth,
+    authenticate_request_bearer_token,
+)
+import importlib
+
+
+class TestAuth:
+    @patch("mlflow_oidc_auth.auth.OAuth")
+    @patch("mlflow_oidc_auth.auth.config")
+    def test_get_oauth_instance(self, mock_config, mock_oauth):
+        mock_app = MagicMock()
+        mock_oauth_instance = MagicMock()
+        mock_oauth.return_value = mock_oauth_instance
+
+        mock_config.OIDC_CLIENT_ID = "mock_client_id"
+        mock_config.OIDC_CLIENT_SECRET = "mock_client_secret"
+        mock_config.OIDC_DISCOVERY_URL = "mock_discovery_url"
+        mock_config.OIDC_SCOPE = "mock_scope"
+
+        result = get_oauth_instance(mock_app)
+
+        mock_oauth.assert_called_once_with(mock_app)
+        mock_oauth_instance.register.assert_called_once_with(
+            name="oidc",
+            client_id="mock_client_id",
+            client_secret="mock_client_secret",
+            server_metadata_url="mock_discovery_url",
+            client_kwargs={"scope": "mock_scope"},
+        )
+        assert result == mock_oauth_instance
+
+    @patch("mlflow_oidc_auth.auth.requests")
+    def test__get_oidc_jwks(self, mock_requests):
+        mock_cache = MagicMock()
+        mock_app = MagicMock()
+        mock_app.logger.debug = MagicMock()
+        mock_requests.get.return_value.json.return_value = {"jwks_uri": "mock_jwks_uri"}
+        mock_cache.get.return_value = None
+
+        # cache and app are imported within the _get_oidc_jwks function
+        mlflow_oidc_app = importlib.import_module("mlflow_oidc_auth.app")
+        with patch.object(mlflow_oidc_app, "cache", mock_cache):
+            with patch.object(mlflow_oidc_app, "app", mock_app):
+                result = _get_oidc_jwks()
+
+                assert len(mock_requests.get.call_args) == 2
+
+                assert mock_requests.get.call_args[0][0] == "mock_jwks_uri"
+                assert mock_requests.get.call_args[1] == {}  # TODO: proper patch for first .get() return_value
+
+                mock_cache.set.assert_called_once_with("jwks", mock_requests.get.return_value.json.return_value, timeout=3600)
+                assert result == mock_requests.get.return_value.json.return_value
+
+    @patch("mlflow_oidc_auth.auth._get_oidc_jwks")
+    @patch("mlflow_oidc_auth.auth.jwt.decode")
+    def test_validate_token(self, mock_jwt_decode, mock_get_oidc_jwks):
+        mock_jwks = {"keys": "mock_keys"}
+        mock_get_oidc_jwks.return_value = mock_jwks
+        mock_payload = MagicMock()
+        mock_jwt_decode.return_value = mock_payload
+
+        token = "mock_token"
+        result = validate_token(token)
+
+        mock_get_oidc_jwks.assert_called_once()
+        mock_jwt_decode.assert_called_once_with(token, mock_jwks)
+        mock_payload.validate.assert_called_once()
+        assert result == mock_payload
+
+    @patch("mlflow_oidc_auth.auth.store")
+    def test_authenticate_request_basic_auth_uses_authenticate_user(self, mock_store):
+        mock_request = MagicMock()
+        mock_request.authorization.username = "mock_username"
+        mock_request.authorization.password = "mock_password"
+        mock_store.authenticate_user.return_value = True
+
+        with patch("mlflow_oidc_auth.auth.request", mock_request):
+            # for some reason decorator doesn't mock flask
+            result = authenticate_request_basic_auth()
+
+            mock_store.authenticate_user.assert_called_once_with("mock_username", "mock_password")
+            assert result == True
+
+    @patch("mlflow_oidc_auth.auth.validate_token")
+    def test_authenticate_request_bearer_token_uses_validate_token(self, mock_validate_token):
+        mock_request = MagicMock()
+        mock_request.authorization.token = "mock_token"
+        mock_validate_token.return_value = MagicMock()
+        with patch("mlflow_oidc_auth.auth.request", mock_request):
+            # for some reason decorator doesn't mock flask
+            result = authenticate_request_bearer_token()
+
+            mock_validate_token.assert_called_once_with("mock_token")
+            assert result == True
+
+    @patch("mlflow_oidc_auth.auth.validate_token")
+    def test_authenticate_request_bearer_token_exception_returns_false(self, mock_validate_token):
+        mock_request = MagicMock()
+        mock_request.authorization.token = "mock_token"
+        mock_validate_token.side_effect = Exception()
+        with patch("mlflow_oidc_auth.auth.request", mock_request):
+            # for some reason decorator doesn't mock flask
+            result = authenticate_request_bearer_token()
+
+            mock_validate_token.assert_called_once_with("mock_token")
+            assert result == False