chore: review fixes

.github/workflows/unit-tests.yml

@@ -6,6 +6,9 @@ on:
       - edited
       - reopened
       - synchronize
+  push:
+    branches:
+      - main
 jobs:
   python-tests:
     name: Run python tests
@@ -30,4 +33,3 @@ jobs:
         uses: SonarSource/sonarcloud-github-action@e44258b109568baa0df60ed515909fc6c72cba92 # v2.3.0
         env:
           SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }}
-          # todo: remove GH App

mlflow_oidc_auth/tests/test_auth.py

@@ -34,28 +34,27 @@ def test_get_oauth_instance(self, mock_config, mock_oauth):
         )
         assert result == mock_oauth_instance
 
-    def test__get_oidc_jwks(self):
-        with patch("mlflow_oidc_auth.auth.requests") as mock_requests:
-            mlflow_oidc_app = importlib.import_module("mlflow_oidc_auth.app")
-            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
-
-            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.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")

mlflow_oidc_auth/tests/test_db_utils.py

@@ -1,8 +1,8 @@
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
-from mlflow_oidc_auth.db.utils import migrate
+from mlflow_oidc_auth.db.utils import migrate, migrate_if_needed
 
 
 class TestMigrate:
@@ -13,3 +13,33 @@ def test_migrate(self, mock_upgrade):
             migrate(engine, "head")
 
         mock_upgrade.assert_called_once()
+
+    @patch("mlflow_oidc_auth.db.utils.MigrationContext")
+    @patch("mlflow_oidc_auth.db.utils.ScriptDirectory")
+    @patch("mlflow_oidc_auth.db.utils.upgrade")
+    def test_migrate_if_needed_not_called_if_not_needed(self, mock_upgrade, mock_script_dir, mock_migration_context):
+        script_dir_mock = MagicMock()
+        script_dir_mock.get_current_head.return_value = "head"
+        mock_script_dir.from_config.return_value = script_dir_mock
+        mock_migration_context.configure.return_value.get_current_revision.return_value = "head"
+
+        engine = create_engine("sqlite:///:memory:")
+        with sessionmaker(bind=engine)():
+            migrate_if_needed(engine, "head")
+
+        mock_upgrade.assert_not_called()
+
+    @patch("mlflow_oidc_auth.db.utils.MigrationContext")
+    @patch("mlflow_oidc_auth.db.utils.ScriptDirectory")
+    @patch("mlflow_oidc_auth.db.utils.upgrade")
+    def test_migrate_if_needed_called_if_needed(self, mock_upgrade, mock_script_dir, mock_migration_context):
+        script_dir_mock = MagicMock()
+        script_dir_mock.get_current_head.return_value = "head"
+        mock_script_dir.from_config.return_value = script_dir_mock
+        mock_migration_context.configure.return_value.get_current_revision.return_value = "not_head"
+
+        engine = create_engine("sqlite:///:memory:")
+        with sessionmaker(bind=engine)():
+            migrate_if_needed(engine, "head")
+
+        mock_upgrade.assert_called_once()

mlflow_oidc_auth/tests/test_sqlalchemy_store.py

@@ -15,52 +15,53 @@ def store(_mock_migrate_if_needed):
 
 
 class TestSqlAlchemyStore:
-    @patch("mlflow_oidc_auth.sqlalchemy_store.SqlAlchemyStore._get_user")
+    @patch(
+        "mlflow_oidc_auth.sqlalchemy_store.SqlAlchemyStore._get_user", return_value=MagicMock(password_hash="hashed_password")
+    )
     @patch("mlflow_oidc_auth.sqlalchemy_store.check_password_hash", return_value=True)
-    def test_authenticate_user(self, _mock_check_password_hash, mock_get_user, store):
-        mock_get_user.return_value = MagicMock(password_hash="hashed_password")
+    def test_authenticate_user(self, mock_check_password_hash, mock_get_user, store):
         auth_result = store.authenticate_user("test_user", "password")
+        mock_check_password_hash.assert_called_once()
         mock_get_user.assert_called_once()
         assert mock_get_user.call_args[0][1] == "test_user"
         assert auth_result is True
 
-    @patch("mlflow_oidc_auth.sqlalchemy_store.SqlAlchemyStore._get_user")
+    @patch(
+        "mlflow_oidc_auth.sqlalchemy_store.SqlAlchemyStore._get_user", return_value=MagicMock(password_hash="hashed_password")
+    )
     @patch("mlflow_oidc_auth.sqlalchemy_store.check_password_hash", return_value=False)
-    def test_authenticate_user_failure(self, _mock_check_password_hash, mock_get_user, store):
-        mock_get_user.return_value = MagicMock(password_hash="hashed_password")
+    def test_authenticate_user_failure(self, mock_check_password_hash, mock_get_user, store):
         auth_result = store.authenticate_user("test_user", "password")
         mock_get_user.assert_called_once()
+        mock_check_password_hash.assert_called_once()
         assert mock_get_user.call_args[0][1] == "test_user"
         assert auth_result is False
 
     @patch("mlflow_oidc_auth.sqlalchemy_store.generate_password_hash", return_value="hashed_password")
-    def test_create_user(self, _generate_password_hash, store):
+    def test_create_user(self, generate_password_hash, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.add = MagicMock()
+        mock_session = MagicMock(flush=MagicMock(), add=MagicMock())
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
 
         user = store.create_user("test_user", "password", "Test User")
 
-        #  display_name="Test User", is_admin=False)
         mock_session.add.assert_called_once()
+        mock_session.flush.assert_called_once()
+        generate_password_hash.assert_called_once_with("password")
+
         assert mock_session.add.call_args[0][0].username == "test_user"
         assert mock_session.add.call_args[0][0].password_hash == "hashed_password"
         assert mock_session.add.call_args[0][0].display_name == "Test User"
         assert mock_session.add.call_args[0][0].is_admin is False
 
-        mock_session.flush.assert_called_once()
         assert user.username == "test_user"
         assert user.display_name == "Test User"
         assert user.is_admin is False
 
     @patch("mlflow_oidc_auth.sqlalchemy_store.generate_password_hash", return_value="hashed_password")
     def test_create_admin_user(self, _generate_password_hash, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.add = MagicMock()
+        mock_session = MagicMock(flush=MagicMock(), add=MagicMock())
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
 
         admin_user = store.create_user("admin_user", "password", "Admin User", is_admin=True)
@@ -73,38 +74,31 @@ def test_create_admin_user(self, _generate_password_hash, store):
 
     def test_create_user_existing(self, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.add = MagicMock(side_effect=IntegrityError("", {}, Exception))
+        mock_session = MagicMock(flush=MagicMock(), add=MagicMock(side_effect=IntegrityError("", {}, Exception)))
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
 
         with pytest.raises(MlflowException):
             store.create_user("test_user", "password", "Test User")
 
     def test_get_user_not_found(self, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.query = MagicMock(side_effect=NoResultFound("", {}, Exception))
+        mock_session = MagicMock(query=MagicMock(side_effect=NoResultFound("", {}, Exception)))
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
 
         with pytest.raises(MlflowException):
             store.get_user("non_existent_user")
 
     @patch("mlflow_oidc_auth.sqlalchemy_store.generate_password_hash", return_value="hashed_password")
     def test_update_user(self, _generate_password_hash, store):
-        retrieved_user = MagicMock()
-        retrieved_user.is_admin = PropertyMock()
-        retrieved_user.password_hash = PropertyMock()
+        retrieved_user = MagicMock(is_admin=PropertyMock(), password_hash=PropertyMock())
         store._get_user = MagicMock(return_value=retrieved_user)
         store.update_user("test_user", password="new_password", is_admin=True)
         assert retrieved_user.is_admin == True
         assert retrieved_user.password_hash == "hashed_password"
 
     def test_delete_user(self, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.delete = MagicMock()
+        mock_session = MagicMock(flush=MagicMock(), delete=MagicMock())
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
         store._get_user = MagicMock(return_value=MagicMock())
 
@@ -119,14 +113,11 @@ def test_create_experiment_permission_validates_permission(self, store):
     @patch("mlflow_oidc_auth.sqlalchemy_store.SqlAlchemyStore._get_user")
     def test_create_experiment_permission(self, mock_get_user, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.add = MagicMock
+        mock_session = MagicMock(flush=MagicMock(), add=MagicMock())
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
 
-        mock_user = MagicMock()
-        mock_user.id = 1
-        store._get_user.return_value = mock_user  # = MagicMock(return_value=mock_user)
+        mock_user = MagicMock(id=1)
+        store._get_user.return_value = mock_user
         mock_get_user.return_value = mock_user
 
         permission = store.create_experiment_permission("1", "test_user", "READ")
@@ -140,9 +131,7 @@ def test_create_registered_model_permission_validates_permission(self, store):
 
     def test_create_registered_model_permission_fails_on_duplicate(self, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.add = MagicMock(side_effect=IntegrityError("", {}, Exception))
+        mock_session = MagicMock(flush=MagicMock(), add=MagicMock(side_effect=IntegrityError("", {}, Exception)))
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
         with pytest.raises(MlflowException):
             store.create_registered_model_permission("model", "test_user", "READ")
@@ -167,9 +156,7 @@ def test_update_registered_model_permission(self, store):
 
     def test_delete_registered_model_permission(self, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.flush = MagicMock()
-        mock_session.delete = MagicMock()
+        mock_session = MagicMock(flush=MagicMock(), delete=MagicMock())
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session
         store._get_registered_model_permission = MagicMock(
             return_value=SqlRegisteredModelPermission(name="model", user_id=1, permission="READ")
@@ -181,8 +168,7 @@ def test_delete_registered_model_permission(self, store):
 
     def test_populate_groups_is_idempotent(self, store):
         store.ManagedSessionMaker = MagicMock()
-        mock_session = MagicMock()
-        mock_session.add = MagicMock()
+        mock_session = MagicMock(add=MagicMock())
         mock_session.query.return_value.filter.return_value.first.return_value = None
         store.ManagedSessionMaker.return_value.__enter__.return_value = mock_session