MINOR: Fix pytests (#21807)

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)
 

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
         """
     )