From 27a0c9e802d3009b17258d1da6cdc8ef2ca2ec96 Mon Sep 17 00:00:00 2001
From: Ayush Shah <ayush@getcollate.io>
Date: Mon, 17 Jul 2023 12:50:11 +0530
Subject: [PATCH] Fix Docker Import (#12455)

---
 docker/__init__.py                            |  0
 docker/validate_compose.py                    | 25 +++++++++++--------
 .../unit/profiler/pandas/test_profiler.py     |  1 -
 3 files changed, 15 insertions(+), 11 deletions(-)
 create mode 100644 docker/__init__.py

diff --git a/docker/__init__.py b/docker/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/docker/validate_compose.py b/docker/validate_compose.py
index 65dd3ed670d..1f135eab825 100644
--- a/docker/validate_compose.py
+++ b/docker/validate_compose.py
@@ -1,10 +1,16 @@
+"""
+Helper functions used for ingestion of sample data into docker by calling airflow dags
+"""
+
+import time
+from pprint import pprint
 from typing import Tuple
 
-from pprint import pprint
 import requests
+from metadata.utils.logger import log_ansi_encoded_string
+from metadata.cli.docker import REQUESTS_TIMEOUT
 from requests.auth import HTTPBasicAuth
-import time
-from metadata.utils.ansi import print_ansi_encoded_string
+
 
 HEADER_JSON = {"Content-Type": "application/json"}
 BASIC_AUTH = HTTPBasicAuth("admin", "admin")
@@ -16,10 +22,10 @@ def get_last_run_info() -> Tuple[str, str]:
     """
     dag_runs = None
     while not dag_runs:
-        print_ansi_encoded_string(message="Waiting for DAG Run data...")
+        log_ansi_encoded_string(message="Waiting for DAG Run data...")
         time.sleep(5)
         runs = requests.get(
-            "http://localhost:8080/api/v1/dags/sample_data/dagRuns", auth=BASIC_AUTH
+            "http://localhost:8080/api/v1/dags/sample_data/dagRuns", auth=BASIC_AUTH, timeout=REQUESTS_TIMEOUT
         ).json()
         dag_runs = runs.get("dag_runs")
 
@@ -33,6 +39,7 @@ def print_last_run_logs() -> None:
     logs = requests.get(
         "http://localhost:8080/api/v1/openmetadata/last_dag_logs?dag_id=sample_data",
         auth=BASIC_AUTH,
+        timeout=REQUESTS_TIMEOUT
     ).text
     pprint(logs)
 
@@ -42,13 +49,11 @@ def main():
     state = None
     while state != "success":
 
-        print_ansi_encoded_string(
-            message="Waiting for sample data ingestion to be a success. We'll show some logs along the way."
+        log_ansi_encoded_string(
+            message="Waiting for sample data ingestion to be a success. We'll show some logs along the way.",
         )
-
         dag_run_id, state = get_last_run_info()
-        print_ansi_encoded_string(message=f"DAG run: [{dag_run_id}, {state}]")
-
+        log_ansi_encoded_string(message=f"DAG run: [{dag_run_id}, {state}]")
         print_last_run_logs()
         time.sleep(10)
 
diff --git a/ingestion/tests/unit/profiler/pandas/test_profiler.py b/ingestion/tests/unit/profiler/pandas/test_profiler.py
index 2f0c1fcb71b..b335c0f0681 100644
--- a/ingestion/tests/unit/profiler/pandas/test_profiler.py
+++ b/ingestion/tests/unit/profiler/pandas/test_profiler.py
@@ -13,7 +13,6 @@
 Test Profiler behavior
 """
 import os
-from concurrent.futures import TimeoutError
 from datetime import datetime
 from unittest import TestCase, mock
 from uuid import uuid4