fix check func issues. update unit tests

src/nebari_plugin_mlflow_aws/__init__.py

@@ -37,6 +37,7 @@ def check(self, stage_outputs: Dict[str, Dict[str, Any]]) -> bool:
         from keycloak import KeycloakAdmin
         from keycloak.exceptions import KeycloakError
 
+        hello = "test"
         # TODO: Module requires EKS cluster is configured for IRSA.  Need to confirm minimum Nebari version once this feature is part of a release.
         # TODO: Also should configure this module to require Nebari version in pyproject.toml?
         try:
@@ -59,10 +60,7 @@ def check(self, stage_outputs: Dict[str, Dict[str, Any]]) -> bool:
             return False
 
         if not self.config.provider == ProviderEnum.aws:
-            print(
-                "\nPlugin 'nebari_plugin_mlflow_aws' developed for AWS only.  Detected provider is {}.".format(self.config.provider)
-            )
-            return False
+            raise KeyError("Plugin 'nebari_plugin_mlflow_aws' developed for 'aws' only.  Detected provider is '{}'.".format(self.config.provider))
 
 
         keycloak_config = self.get_keycloak_config(stage_outputs)
@@ -129,7 +127,7 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]):
             cluster_oidc_issuer_url = stage_outputs["stages/02-infrastructure"]["cluster_oidc_issuer_url"]["value"]
 
         except KeyError:
-            raise Exception("Prerequisite stage output(s) not found: stages/04-kubernetes-ingress")
+            raise Exception("Prerequisite stage output(s) not found: stages/02-infrastructure, stages/04-kubernetes-ingress")
 
         chart_ns = self.config.mlflow.namespace
         create_ns = True

tests/unit/test_plugin.py

@@ -0,0 +1,97 @@
+import pytest
+from nebari_plugin_mlflow_aws import MlflowStage, MlflowConfig, InputSchema
+
+class TestConfig(InputSchema):
+    __test__ = False
+    namespace: str
+    domain: str
+    escaped_project_name: str = ""
+    provider: str
+
+@pytest.fixture(autouse=True)
+def mock_keycloak_connection(monkeypatch):
+    monkeypatch.setattr("nebari_plugin_mlflow_aws.MlflowStage.check._attempt_keycloak_connection", lambda: True)
+
+def test_ctor():
+    sut = MlflowStage(output_directory = None, config = None)
+    assert sut.name == "mlflow"
+    assert sut.priority == 102
+
+def test_input_vars():
+    config = TestConfig(namespace = "nebari-ns", domain = "my-test-domain.com", escaped_project_name="testprojectname", provider="aws")
+    sut = MlflowStage(output_directory = None, config = config)
+
+
+    stage_outputs = get_stage_outputs()
+    #sut.check(stage_outputs)
+    result = sut.input_vars(stage_outputs)
+    assert result["chart_name"] == "mlflow"
+    assert result["project_name"] == "testprojectname"
+    assert result["realm_id"] == "test-realm"
+    assert result["client_id"] == "mlflow"
+    assert result["base_url"] == "https://my-test-domain.com/mlflow"
+    assert result["external_url"] == "https://my-test-domain.com/auth/"
+    assert result["valid_redirect_uris"] == ["https://my-test-domain.com/mlflow/_oauth"]
+    assert result["signing_key_ref"] == {'kind': 'Deployment', 'name': 'forwardauth-deployment', 'namespace': 'nebari-ns'}
+    assert result["create_namespace"] == False
+    assert result["namespace"] == "nebari-ns"
+    assert result["ingress_host"] == "my-test-domain.com"
+    assert result["cluster_oidc_issuer_url"] == "https://test-oidc-url.com"
+    assert result["overrides"] == {}
+
+#def test_incompatible_cloud():
+# TODO
+#
+#def test_default_namespace():
+#    config = TestConfig(namespace = "nebari-ns", domain = "my-test-domain.com")
+#    sut = MlflowStage(output_directory = None, config = config)
+#
+#    stage_outputs = get_stage_outputs()
+#    result = sut.input_vars(stage_outputs)
+#    assert result["create_namespace"] == False
+#    assert result["namespace"] == "nebari-ns"
+#
+#def test_chart_namespace():
+#    config = TestConfig(namespace = "nebari-ns", domain = "my-test-domain.com", label_studio = MlflowStage(namespace = "label-studio-ns"))
+#    sut = MlflowStage(output_directory = None, config = config)
+#
+#    stage_outputs = get_stage_outputs()
+#    result = sut.input_vars(stage_outputs)
+#    assert result["create_namespace"] == True
+#    assert result["namespace"] == "label-studio-ns"
+#
+#def test_chart_overrides():
+#    config = TestConfig(namespace = "nebari-ns", domain = "my-test-domain.com", label_studio = MlflowStage(values = { "foo": "bar" }))
+#    sut = MlflowStage(output_directory = None, config = config)
+#
+#    stage_outputs = get_stage_outputs()
+#    result = sut.input_vars(stage_outputs)
+#    assert result["overrides"] == { "foo": "bar" }
+#
+def get_stage_outputs():
+    return {
+        "stages/02-infrastructure": {
+            "cluster_oidc_issuer_url": {
+                "value": "https://test-oidc-url.com"
+            }
+        },
+        "stages/04-kubernetes-ingress": {
+            "domain": "my-test-domain.com"
+        },
+        "stages/05-kubernetes-keycloak": {
+            "keycloak_credentials": {
+                "value": {
+                    "url": "https://my-test-domain.com",
+                    "username": "testuser",
+                    "password": "testpassword",
+                    "realm": "testmasterrealm",
+                    "client_id": "testmasterclientid"
+                }
+            }
+        },
+        "stages/06-kubernetes-keycloak-configuration": {
+            "realm_id": {
+                "value": "test-realm"
+            }
+        }
+    }