From ea285923932693264f8d0828cc821438ad5582a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bar=C4=B1=C5=9F=20Can=20Durak?=
 <36421093+bcdurak@users.noreply.github.com>
Date: Thu, 19 Sep 2024 16:21:15 +0200
Subject: [PATCH] MLFlow integration requirements duplicate fix (#3011)

* added copy

* removed requirements

* fixed integration test
---
 src/zenml/integrations/mlflow/__init__.py     | 29 +++++++++----------
 .../functional/cli/test_integration.py        |  6 ++--
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/src/zenml/integrations/mlflow/__init__.py b/src/zenml/integrations/mlflow/__init__.py
index 6cd1ef25eff..fcd9c1f0545 100644
--- a/src/zenml/integrations/mlflow/__init__.py
+++ b/src/zenml/integrations/mlflow/__init__.py
@@ -37,20 +37,6 @@ class MlflowIntegration(Integration):
 
     NAME = MLFLOW
 
-    REQUIREMENTS = [
-        "mlflow>=2.1.1,<3",
-        # TODO: remove this requirement once rapidjson is fixed
-        "python-rapidjson<1.15",
-        # When you do:
-        # pip install zenml
-        # You get all our required dependencies. However, if you follow it with:
-        # zenml integration install mlflow
-        # This downgrades pydantic to v1 even though mlflow does not have
-        # any issues with v2. This is why we have to pin it here so a downgrade
-        # will not happen.
-        "pydantic>=2.8.0,<2.9.0",
-    ]
-
     REQUIREMENTS_IGNORED_ON_UNINSTALL = [
         "python-rapidjson",
         "pydantic",
@@ -71,7 +57,20 @@ def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
         from zenml.integrations.numpy import NumpyIntegration
         from zenml.integrations.pandas import PandasIntegration
 
-        reqs = cls.REQUIREMENTS
+        reqs = [
+            "mlflow>=2.1.1,<3",
+            # TODO: remove this requirement once rapidjson is fixed
+            "python-rapidjson<1.15",
+            # When you do:
+            #       pip install zenml
+            # You get all our required dependencies. However, if you follow it
+            # with:
+            #       zenml integration install mlflow
+            # This downgrades pydantic to v1 even though mlflow does not have
+            # any issues with v2. This is why we have to pin it here so a
+            # downgrade will not happen.
+            "pydantic>=2.8.0,<2.9.0",
+        ]
         if sys.version_info.minor >= 12:
             logger.debug(
                 "The MLflow integration on Python 3.12 and above is not yet "
diff --git a/tests/integration/functional/cli/test_integration.py b/tests/integration/functional/cli/test_integration.py
index 14b852f133e..3b9e511acf7 100644
--- a/tests/integration/functional/cli/test_integration.py
+++ b/tests/integration/functional/cli/test_integration.py
@@ -229,9 +229,9 @@ def test_integration_requirements_exporting(tmp_path) -> None:
     from zenml.integrations.mlflow import MlflowIntegration
 
     flow_integration_requirements = set(
-        AirflowIntegration.REQUIREMENTS
-        + KubeflowIntegration.REQUIREMENTS
-        + MlflowIntegration.REQUIREMENTS
+        AirflowIntegration.get_requirements()
+        + KubeflowIntegration.get_requirements()
+        + MlflowIntegration.get_requirements()
     )
 
     command = [