Repository: Galileo-Galilei/kedro-mlflow (language: python)

Issue:
Increase test speed
- [ ] Some tests seem very slow, I should do some profiling and make them faster for developer experience. Use ``pytest --durations 0`` to get all tests durations. 
- [ ] somes tests create empty ``mlruns`` folder anywhere, I should get rid of this


Discussion / hints:
Almost all slow tests are due to call to ``mlflow.log_model``
![image](https://github.com/Galileo-Galilei/kedro-mlflow/assets/29451317/6003c3bb-76ba-40ae-bc95-0a895a90c7b4)

There is a known bug when using mlflow locally on windows 10 :
https://github.com/mlflow/mlflow/issues/4062

On windows 11, I still have the issue with the follwoing code: 

```bash
mlflow server

INFO:waitress:Serving on http://127.0.0.1:5000
```

```python
import time
from pathlib import Path

import mlflow
from sklearn.tree import DecisionTreeClassifier

tracking_uri = "http://127.0.0.1:5000"
mlflow.set_tracking_uri(tracking_uri)
model = DecisionTreeClassifier()
st = time.time()
with mlflow.start_run():
    mlflow.sklearn.log_model(model, artifact_path="demo_model")
time.time() - st

# > always get ~2.5 s

```

Fix patch:
diff --git a/.gitignore b/.gitignore
index e7a88bb2..282c2f07 100644
--- a/.gitignore
+++ b/.gitignore
@@ -134,6 +134,9 @@ dmypy.json
 # mlflow
 mlruns/
 
+# ruff
+.ruff_cache
+
 debug/
 *.xlsx
 *.pptx


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 700d88be..d7882b42 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -6,6 +6,7 @@
 from cookiecutter.main import cookiecutter
 from kedro import __version__ as kedro_version
 from kedro.framework.cli.starters import TEMPLATE_PATH
+from mlflow import MlflowClient
 
 from kedro_mlflow.framework.cli.cli import TEMPLATE_FOLDER_PATH
 from kedro_mlflow.framework.cli.cli_utils import write_jinja_template
@@ -13,19 +14,39 @@
 _FAKE_PROJECT_NAME = "fake_project"
 
 
+@pytest.fixture
+def tracking_uri(tmp_path):
+    tracking_uri = (tmp_path / "mlruns").as_uri()
+    return tracking_uri
+
+
+@pytest.fixture
+def mlflow_client(tracking_uri):
+    mlflow.set_tracking_uri(tracking_uri)
+    client = MlflowClient(tracking_uri)
+    return client
+
+
 @pytest.fixture(autouse=True)
 def cleanup_mlflow_after_runs():
     # A test function will be run at this point
     yield
     while mlflow.active_run():
         mlflow.end_run()
-    mlflow.set_experiment("Default")
+
+    import mlflow.tracking.fluent as mtf
+
+    # if set_experiment has been called before, it stores the experiment_id
+    # as a global variable, so if we change the tracking_uri afterwards
+    # mlflow is completly lost because the experiment id no longer exists
+    # we just reset it after a test, like in a brand new session
+    if hasattr(mtf, "_active_experiment_id"):
+        mtf._active_experiment_id = None
 
     if "MLFLOW_TRACKING_URI" in os.environ:
         os.environ.pop("MLFLOW_TRACKING_URI")
 
-
-# see https://github.com/kedro-org/kedro/blob/859f98217eed12208a922b771a97cbfb82ba7e80/tests/framework/session/test_session.py#L173
+    # see https://github.com/kedro-org/kedro/blob/859f98217eed12208a922b771a97cbfb82ba7e80/tests/framework/session/test_session.py#L173
 
 
 @pytest.fixture
diff --git a/tests/io/artifacts/test_mlflow_artifact_dataset.py b/tests/io/artifacts/test_mlflow_artifact_dataset.py
index 0a1ac3f7..ec3580e2 100644
--- a/tests/io/artifacts/test_mlflow_artifact_dataset.py
+++ b/tests/io/artifacts/test_mlflow_artifact_dataset.py
@@ -6,17 +6,11 @@
 from kedro.io import PartitionedDataset
 from kedro_datasets.pandas import CSVDataset
 from kedro_datasets.pickle import PickleDataset
-from mlflow.tracking import MlflowClient
 from pytest_lazyfixture import lazy_fixture
 
 from kedro_mlflow.io.artifacts import MlflowArtifactDataset
 
 
-@pytest.fixture
-def tracking_uri(tmp_path):
-    return tmp_path / "mlruns"
-
-
 @pytest.fixture
 def df1():
     return pd.DataFrame({"col1": [1, 2, 3], "col2": [4, 5, 6]})
@@ -46,10 +40,8 @@ def df2():
     ],
 )
 def test_mlflow_csv_pkl_dataset_save_reload(
-    tmp_path, tracking_uri, dataset, extension, data, artifact_path
+    tmp_path, mlflow_client, dataset, extension, data, artifact_path
 ):
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
     filepath = (tmp_path / "data").with_suffix(extension)
 
     mlflow_dataset = MlflowArtifactDataset(
@@ -80,10 +72,8 @@ def test_mlflow_csv_pkl_dataset_save_reload(
     [(False), (True)],
 )
 def test_artifact_dataset_save_with_run_id(
-    tmp_path, tracking_uri, df1, exists_active_run
+    tmp_path, mlflow_client, df1, exists_active_run
 ):
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
     nb_runs = 0
     # create a first run and get its id
     with mlflow.start_run():
@@ -123,13 +113,11 @@ def test_artifact_dataset_save_with_run_id(
         mlflow.end_run()
 
 
-def test_is_versioned_dataset_logged_correctly_in_mlflow(tmp_path, tracking_uri, df1):
+def test_is_versioned_dataset_logged_correctly_in_mlflow(tmp_path, mlflow_client, df1):
     """Check if versioned dataset is logged correctly in MLflow as artifact.
 
     For versioned datasets just artifacts from current run should be logged.
     """
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
 
     with mlflow.start_run():
         run_id = mlflow.active_run().info.run_id
@@ -171,14 +159,11 @@ def test_is_versioned_dataset_logged_correctly_in_mlflow(tmp_path, tracking_uri,
         assert df1.equals(mlflow_csv_dataset.load())  # and must loadable
 
 
-def test_artifact_dataset_logging_deactivation(tmp_path, tracking_uri):
+def test_artifact_dataset_logging_deactivation(tmp_path, mlflow_client):
     mlflow_pkl_dataset = MlflowArtifactDataset(
         dataset=dict(type=PickleDataset, filepath=(tmp_path / "df1.csv").as_posix())
     )
 
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
-
     mlflow_pkl_dataset._logging_activated = False
 
     all_runs_id_beginning = {
@@ -208,7 +193,7 @@ def test_mlflow_artifact_logging_deactivation_is_bool(tmp_path):
 
 
 def test_artifact_dataset_load_with_run_id(tmp_path, tracking_uri, df1, df2):
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
+    mlflow.set_tracking_uri(tracking_uri)
 
     # define the logger
     mlflow_csv_dataset = MlflowArtifactDataset(
@@ -236,7 +221,7 @@ def test_artifact_dataset_load_with_run_id(tmp_path, tracking_uri, df1, df2):
 def test_artifact_dataset_load_with_run_id_and_artifact_path(
     tmp_path, tracking_uri, df1, artifact_path
 ):
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
+    mlflow.set_tracking_uri(tracking_uri)
 
     # save first and retrieve run id
     mlflow_csv_dataset1 = MlflowArtifactDataset(
@@ -264,11 +249,8 @@ def test_artifact_dataset_load_with_run_id_and_artifact_path(
 
 @pytest.mark.parametrize("artifact_path", [None, "partitioned_data"])
 def test_partitioned_dataset_save_and_reload(
-    tmp_path, tracking_uri, artifact_path, df1, df2
+    tmp_path, mlflow_client, artifact_path, df1, df2
 ):
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
-
     mlflow_dataset = MlflowArtifactDataset(
         artifact_path=artifact_path,
         dataset=dict(
diff --git a/tests/io/metrics/test_mlflow_metrics_dataset.py b/tests/io/metrics/test_mlflow_metrics_dataset.py
index 2ebece9f..d177af99 100644
--- a/tests/io/metrics/test_mlflow_metrics_dataset.py
+++ b/tests/io/metrics/test_mlflow_metrics_dataset.py
@@ -38,11 +38,6 @@ def assert_are_metrics_logged(
     assert True
 
 
-@pytest.fixture
-def tracking_uri(tmp_path):
-    return tmp_path / "mlruns"
-
-
 @pytest.fixture
 def metrics():
     return {
@@ -80,12 +75,11 @@ def metrics3():
         (lazy_fixture("metrics3"), "test"),
     ],
 )
-def test_mlflow_metrics_dataset_saved_and_logged(tmp_path, tracking_uri, data, prefix):
+def test_mlflow_metrics_dataset_saved_and_logged(mlflow_client, data, prefix):
     """Check if MlflowMetricsHistoryDataset can be saved in catalog when filepath is given,
     and if logged in mlflow.
     """
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
+
     mlflow_metrics_dataset = MlflowMetricsHistoryDataset(prefix=prefix)
 
     with mlflow.start_run():
@@ -108,14 +102,12 @@ def test_mlflow_metrics_dataset_saved_and_logged(tmp_path, tracking_uri, data, p
         assert data[data_key] == catalog_metrics[k]
 
 
-def test_mlflow_metrics_dataset_saved_without_run_id(tmp_path, tracking_uri, metrics3):
+def test_mlflow_metrics_dataset_saved_without_run_id(mlflow_client, metrics3):
     """Check if MlflowMetricsHistoryDataset can be saved in catalog when filepath is given,
     and if logged in mlflow.
     """
     prefix = "test_metric"
 
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
     mlflow_metrics_dataset = MlflowMetricsHistoryDataset(prefix=prefix)
 
     # a mlflow run_id is automatically created
@@ -125,13 +117,13 @@ def test_mlflow_metrics_dataset_saved_without_run_id(tmp_path, tracking_uri, met
     assert_are_metrics_logged(metrics3, mlflow_client, run_id, prefix)
 
 
-def test_mlflow_metrics_dataset_exists(tmp_path, tracking_uri, metrics3):
+def test_mlflow_metrics_dataset_exists(tracking_uri, metrics3):
     """Check if MlflowMetricsHistoryDataset is well identified as
     existing if it has already been saved.
     """
     prefix = "test_metric"
 
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
+    mlflow.set_tracking_uri(tracking_uri)
     mlflow_metrics_dataset = MlflowMetricsHistoryDataset(prefix=prefix)
 
     # a mlflow run_id is automatically created
@@ -139,13 +131,12 @@ def test_mlflow_metrics_dataset_exists(tmp_path, tracking_uri, metrics3):
     assert mlflow_metrics_dataset.exists()
 
 
-def test_mlflow_metrics_dataset_does_not_exist(tmp_path, tracking_uri, metrics3):
+def test_mlflow_metrics_dataset_does_not_exist(tracking_uri):
     """Check if MlflowMetricsHistoryDataset is well identified as
     not existingif it has never been saved.
     """
-
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow.start_run()  # starts a run toenable mlflow_metrics_dataset to know where to seacrh
+    mlflow.set_tracking_uri(tracking_uri)
+    mlflow.start_run()  # starts a run to enable mlflow_metrics_dataset to know where to seacrh
     run_id = mlflow.active_run().info.run_id
     mlflow.end_run()
     mlflow_metrics_dataset = MlflowMetricsHistoryDataset(
@@ -155,14 +146,12 @@ def test_mlflow_metrics_dataset_does_not_exist(tmp_path, tracking_uri, metrics3)
     assert not mlflow_metrics_dataset.exists()
 
 
-def test_mlflow_metrics_dataset_fails_with_invalid_metric(
-    tmp_path, tracking_uri, metrics3
-):
+def test_mlflow_metrics_dataset_fails_with_invalid_metric(tracking_uri):
     """Check if MlflowMetricsHistoryDataset is well identified as
     not existingif it has never been saved.
     """
 
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
+    mlflow.set_tracking_uri(tracking_uri)
     mlflow_metrics_dataset = MlflowMetricsHistoryDataset(prefix="test_metric")
 
     with pytest.raises(
@@ -173,12 +162,9 @@ def test_mlflow_metrics_dataset_fails_with_invalid_metric(
         )  # key: value is not valid, you must specify {key: {value, step}}
 
 
-def test_mlflow_metrics_logging_deactivation(tracking_uri, metrics):
+def test_mlflow_metrics_logging_deactivation(mlflow_client, metrics):
     mlflow_metrics_dataset = MlflowMetricsHistoryDataset(prefix="hello")
 
-    mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
-
     mlflow_metrics_dataset._logging_activated = False
 
     all_runs_id_beginning = {
diff --git a/tests/io/models/test_mlflow_model_tracking_dataset.py b/tests/io/models/test_mlflow_model_tracking_dataset.py
index 97b81c78..4f5b188e 100644
--- a/tests/io/models/test_mlflow_model_tracking_dataset.py
+++ b/tests/io/models/test_mlflow_model_tracking_dataset.py
@@ -7,7 +7,6 @@
 from kedro.io.core import DatasetError
 from kedro.pipeline import Pipeline, node
 from kedro_datasets.pickle import PickleDataset
-from mlflow.tracking import MlflowClient
 from sklearn.linear_model import LinearRegression
 
 from kedro_mlflow.io.models import MlflowModelTrackingDataset
@@ -25,24 +24,12 @@ def linreg_model():
     return linreg_model
 
 
-@pytest.fixture
-def tracking_uri(tmp_path):
-    tracking_uri = (tmp_path / "mlruns").as_uri()
-    return tracking_uri
-
-
 @pytest.fixture
 def tmp_folder():
     tmp_folder = TemporaryDirectory()
     return tmp_folder
 
 
-@pytest.fixture
-def mlflow_client(tracking_uri):
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri)
-    return mlflow_client
-
-
 @pytest.fixture
 def prexisting_run_id(tracking_uri):
     mlflow.set_tracking_uri(tracking_uri)
@@ -172,16 +159,10 @@ def test_save_sklearn_flavor_with_run_id_and_already_active_run(tracking_uri):
 
 @pytest.mark.parametrize("active_run_when_loading", [False, True])
 def test_save_and_load_sklearn_flavor_with_run_id(
-    tracking_uri, mlflow_client, linreg_model, active_run_when_loading
+    mlflow_client, linreg_model, active_run_when_loading
 ):
-    mlflow.set_tracking_uri(tracking_uri)
-    # close all opened mlflow runs to avoid interference between tests
-    while mlflow.active_run():
-        mlflow.end_run()
-
-    mlflow.start_run()
-    existing_run_id = mlflow.active_run().info.run_id
-    mlflow.end_run()
+    with mlflow.start_run():
+        existing_run_id = mlflow.active_run().info.run_id
 
     artifact_path = "my_linreg"
     model_config = {
@@ -213,13 +194,8 @@ def test_save_and_load_sklearn_flavor_with_run_id(
 
 @pytest.mark.parametrize("initial_active_run", [False, True])
 def test_save_and_load_sklearn_flavor_without_run_id(
-    tracking_uri, mlflow_client, linreg_model, initial_active_run
+    mlflow_client, linreg_model, initial_active_run
 ):
-    mlflow.set_tracking_uri(tracking_uri)
-    # close all opened mlflow runs to avoid interference between tests
-    while mlflow.active_run():
-        mlflow.end_run()
-
     artifact_path = "my_linreg"
     model_config = {
         "name": "linreg",
@@ -343,6 +319,7 @@ def test_pyfunc_flavor_python_model_save_and_load(
 
 
 def test_pyfunc_flavor_wrong_pyfunc_workflow(tracking_uri):
+    mlflow.set_tracking_uri(tracking_uri)
     model_config = {
         "name": "kedro_pipeline_model",
         "config": {
@@ -359,12 +336,9 @@ def test_pyfunc_flavor_wrong_pyfunc_workflow(tracking_uri):
         MlflowModelTrackingDataset.from_config(**model_config)
 
 
-def test_mlflow_model_tracking_logging_deactivation(tracking_uri, linreg_model):
+def test_mlflow_model_tracking_logging_deactivation(mlflow_client, linreg_model):
     mlflow_model_tracking_dataset = MlflowModelTrackingDataset(flavor="mlflow.sklearn")
 
-    mlflow.set_tracking_uri(tracking_uri)
-    mlflow_client = MlflowClient(tracking_uri=tracking_uri)
-
     mlflow_model_tracking_dataset._logging_activated = False
 
     all_runs_id_beginning = {