test: add mlflow integration tests

Author: Timon Viola

.containerignore

@@ -0,0 +1,2 @@
+!share
+*

.pre-commit-config.yaml

@@ -15,13 +15,15 @@ repos:
     - id: typos
       stages: [commit]
       exclude: "^typings/.*"
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.10.1  # Use the sha / tag you want to point at
+  - repo: local
     hooks:
-    -   id: mypy
-        language: system
-        pass_filenames: false
-        args: ['.']
+    - id: mypy-hatch
+      name: mypy-hatch
+      language: system
+      pass_filenames: false
+      args: ['.']
+      entry: hatch
+      args: ["run", "types:check"]
   - repo: https://github.com/compilerla/conventional-pre-commit
     rev: v3.3.0
     hooks:

docker-compose.yaml

@@ -9,12 +9,11 @@ version: '3.8'
 services:
   airflow:
     build: .
+    #image: pytest74406-airflow:latest
     ports:
       - "8080:8080"
     container_name: airflow
     restart: always
-    depends_on:
-      - db
     volumes:
       - ./tests/dags:/opt/airflow/dags
       - ./:/opt/dagcellent/

pyproject.toml

@@ -1,10 +1,10 @@
 [build-system]
-requires = ["hatchling", "versioningit"]
+requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [project]
 name = "dagcellent"
-dynamic = ["version"]
+version = "0.2.0"
 description = ''
 readme = "README.md"
 requires-python = ">=3.11"
@@ -28,7 +28,7 @@ classifiers = [
 dependencies = [
     "apache-airflow>=2.9.1",
     "apache-airflow-providers-amazon>=7.2.0",
-    "apache-airflow-providers-microsoft-mssql >=3.2.0",
+    "apache-airflow-providers-microsoft-mssql >=4.0.0",
     "apache-airflow-providers-snowflake >= 5.1",
     "apache-airflow-providers-common-sql[pandas,openlineage]",
     "tomli >= 2.0.1",
@@ -64,18 +64,17 @@ extra-dependencies = [
     "pre-commit == 3.7.*",
     "ruff == 0.4.4",
     "mypy == 1.10.*",
-    "versioningit == 3.1.*",
     "towncrier == 23.11.*",
 ]
 
 [tool.hatch.envs.dev.scripts]
 install = "pre-commit install --hook-type commit-msg"
-version = "versioningit"
 changelog = "git-cliff -o CHANGELOG.md"
 
 [tool.hatch.envs.types]
 extra-dependencies = [
   "mypy>=1.0.0",
+  "types-requests",
 ]
 [tool.hatch.envs.types.scripts]
 check = "mypy --install-types --non-interactive {args:src/dagcellent tests}"
@@ -96,16 +95,15 @@ extra-dependencies = [
     "pytest >= 8.0.0",
     "pytest-cov",
     "pytest-mock >= 3.14.0",
+    "pytest-docker",
 ]
 
 [[tool.hatch.envs.test.matrix]]
 python = ["3.11", "3.12"]
 
 [tool.hatch.envs.test.scripts]
-test = "pytest --cov=dagcellent --cov-report=term-missing --cov-report=xml --cov-report=html tests"
+test = "pytest --cov=dagcellent --cov-report=term-missing --cov-report=xml --cov-report=html --container-scope module tests"
 
-[tool.hatch.version]
-source = "versioningit"
 
 [tool.coverage.run]
 source_pkgs = ["dagcellent", "tests"]
@@ -311,24 +309,4 @@ exclude = [
     "^typings/*"
 ]
 
-[tool.versioningit]
-default-version = "0.0.0-unknown"
-
-[tool.versioningit.vcs]
-# The method key:
-method = "git"  # <- The method name
-
-# Parameters to pass to the method:
-match = ["v*"]
-default-tag = "0.0.1"
-
-[tool.hatch.build.hooks.versioningit-onbuild]
-source-file = "src/dagcellent/__init__.py"
-build-file = "dagcellent/__init__.py"
-require-match = true
-
-[tool.versioningit.format]
-distance = "{next_version}.dev{distance}+{vcs}{rev}"
-dirty = "{version}+dirty"
-distance-dirty = "{next_version}.dev{distance}+{vcs}{rev}.dirty"
 

src/dagcellent/operators/mlflow/hooks.py

@@ -30,7 +30,7 @@
 def _mlflow_request_wrapper(query: Callable[P, T]) -> T:
     """Wrap requests around MlflowException."""
     try:
-        res = query()
+        res = query()  # type: ignore[call-arg]
     except mlflow.MlflowException as exc:
         _msg = "Error during mlflow query."
         _LOGGER.error(_msg, exc_info=exc)
@@ -63,11 +63,17 @@ def get_latest_model_version(
 
         Returns:
             dict: hashmap with version and run_id of latest model
+
+        Raises:
+            ValueError: No model returned with given name.
         """
         query = functools.partial(
             self.client.search_model_versions, f"name = '{model_name}'"
         )
         model_reigstry_info = _mlflow_request_wrapper(query)
+        if len(model_reigstry_info) == 0:
+            _msg = f"No models found with name {model_name}."
+            raise ValueError(_msg)
         latest_version = functools.reduce(
             lambda x, y: x if int(x.version) > int(y.version) else y,
             model_reigstry_info,
@@ -110,6 +116,11 @@ def transition_model_version_stage(
         Returns:
             mlflow.entities.model_registry.ModelVersion: model version
         """
+        warn(
+            "This function is deprecated in mlflow 2.9.0",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         query = functools.partial(
             self.client.transition_model_version_stage,
             name,
@@ -135,7 +146,11 @@ def get_latest_versions(
         Returns:
             list[mlflow.entities.model_registry.ModelVersion]: list of model versions
         """
-        warn("This is deprecated in mlflow 2.9.0", DeprecationWarning, stacklevel=2)
+        warn(
+            "This function is deprecated in mlflow 2.9.0",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         query = functools.partial(self.client.get_latest_versions, name, stages)
         return _mlflow_request_wrapper(query)
 

tests/integration/docker-compose.override.mssql.yaml

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from pathlib import Path
+from warnings import warn
+
+import pytest
+import requests
+from requests.exceptions import ConnectionError
+
+
+def _is_responsive(url: str):
+    try:
+        response = requests.get(url)
+        if response.status_code == 200:
+            return True
+    except ConnectionError:
+        return False
+
+
+def _remove_path(path: Path):
+    if path.is_file() or path.is_symlink():
+        path.unlink()
+        return
+    for p in path.iterdir():
+        _remove_path(p)
+    path.rmdir()
+
+
+@pytest.fixture(scope="session")
+def docker_compose_file(pytestconfig: pytest.Config):
+    return [
+        str(pytestconfig.rootpath / "docker-compose.yaml"),
+        str(Path(__file__).parent / "docker-compose.override.mlflow.yaml"),
+    ]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def mlflow_service(docker_ip, docker_services, pytestconfig: pytest.Config):
+    """Ensure that HTTP service is up and responsive."""
+
+    # `port_for` takes a container port and returns the corresponding host port
+    port = docker_services.port_for("mlflow", 5000)
+    url = f"http://{docker_ip}:{port}"
+    docker_services.wait_until_responsive(
+        timeout=30.0, pause=0.1, check=lambda: _is_responsive(url)
+    )
+    yield url
+    # after run, remove the root 'mlruns' folder
+    _p = pytestconfig.rootpath / "mlruns"
+    if not _p.exists():
+        return
+
+    if _p.is_dir():
+        _remove_path(_p)
+        return
+    warn(
+        "Could not clean up `mlruns` artifact. This might cause unexpected"
+        "behaviour in the next test suite execution.",
+        stacklevel=2,
+    )
+
+
+@pytest.fixture(scope="module")
+def mlflow_hook(mlflow_service: str):
+    """Get MLFlow client wrapper."""
+    from dagcellent.operators.mlflow.hooks import MlflowHook
+
+    return MlflowHook(mlflow_service)

tests/integration/docker-compose.override.psql.yaml

@@ -0,0 +1,39 @@
+version: '3.8'
+
+services:
+  #  minio:
+  #    image: minio/minio:latest
+  #    ports:
+  #     - "9000:9000"
+  #    environment:
+  #      - "MINIO_ACCESS_KEY=mlflow_access_key"
+  #      - "MINIO_SECRET_KEY=mlflow_secret_key"
+  #    volumes:
+  #      - minio_data:/data
+  #    command: server /data
+  #
+  mlflow:
+    image: ghcr.io/mlflow/mlflow:v2.20.3
+    #    build:
+     # context: ..
+      # dockerfile: ghcr.io/mlflow/mlflow:v2.20.3
+      #      args:
+      #  - MLFLOW_VERSION=">=2.2"
+    container_name: mlflow
+    restart: on-failure
+    ports:
+      - "5005:5000"
+    #depends_on:
+      #  - postgres
+    environment:
+      - GUNICORN_CMD_ARGS="--bind=0.0.0.0"
+    #  - MLFLOW_S3_ENDPOINT_URL=http://minio:9000
+    command: mlflow server --backend-store-uri=sqlite:///mlruns.db --default-artifact-root=file:mlruns --host 0.0.0.0 --port 5000
+    #command: mlflow ui --host 0.0.0.0 --port 8888
+    # command: mlflow server --backend-store-uri postgresql://mlflow:mlflow@postgres:5432/mlflow --default-artifact-root s3://mlflow/ --host 0.0.0.0
+
+#volumes:
+#  postgres-db-volume:
+#  logs:
+#  minio_data:
+