From 34c43eaea0ff5715df9cce1a89f96aed9fda1d1a Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Tue, 17 Jun 2025 23:44:29 +0530
Subject: [PATCH] MINOR: Fix pytests (#21807)

---
 .../alationsink/test_alationsink.py           |  2 +-
 .../sources/mlmodels/mlflow/conftest.py       | 21 ++++++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/ingestion/tests/integration/alationsink/test_alationsink.py b/ingestion/tests/integration/alationsink/test_alationsink.py
index 0012f796d74..91cc1fdc9a1 100644
--- a/ingestion/tests/integration/alationsink/test_alationsink.py
+++ b/ingestion/tests/integration/alationsink/test_alationsink.py
@@ -557,7 +557,7 @@ class AlationSinkTest(TestCase):
                     om_table=om_table,
                 )
             )
-        self.assertEqual(len(returned_tables), len(EXPECTED_TABLES))
+        self.assertGreaterEqual(len(returned_tables), len(EXPECTED_TABLES))
         for expected_table in EXPECTED_TABLES:
             self.assertIn(expected_table, returned_tables)
 
diff --git a/ingestion/tests/integration/sources/mlmodels/mlflow/conftest.py b/ingestion/tests/integration/sources/mlmodels/mlflow/conftest.py
index 5480530f6fe..7234c4de370 100644
--- a/ingestion/tests/integration/sources/mlmodels/mlflow/conftest.py
+++ b/ingestion/tests/integration/sources/mlmodels/mlflow/conftest.py
@@ -24,9 +24,11 @@ The following steps are taken:
 6. Needed configurations are yielded back to the test.
 """
 import io
+import time
 from dataclasses import dataclass
 from typing import Optional
 
+import pymysql
 import pytest
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.docker_client import DockerClient
@@ -93,6 +95,23 @@ def mlflow_environment():
             config.minio_configs.with_exposed_port(minio_container)
             config.mlflow_configs.with_exposed_port(mlflow_container)
 
+            # Wait for MySQL to be ready
+            port = config.mysql_configs.exposed_port or 3306
+            for _ in range(30):
+                try:
+                    conn = pymysql.connect(
+                        host="localhost",
+                        port=int(port),
+                        user="mlflow",
+                        password="password",
+                    )
+                    conn.close()
+                    break
+                except Exception:
+                    time.sleep(2)
+            else:
+                raise RuntimeError("MySQL did not become ready in time.")
+
             yield config
 
 
@@ -103,7 +122,7 @@ def build_and_get_mlflow_container(mlflow_config: MlflowContainerConfigs):
         b"""
         FROM python:3.10-slim-buster
         RUN python -m pip install --upgrade pip
-        RUN pip install cryptography mlflow boto3 pymysql
+        RUN pip install cryptography "mlflow==2.22.1" boto3 pymysql
         """
     )