From 0e418b527e64b9314c2a4da1df7794b129ac21cb Mon Sep 17 00:00:00 2001
From: Harshal Sheth <hsheth2@gmail.com>
Date: Thu, 1 Feb 2024 16:33:15 -0800
Subject: [PATCH] fix(ingest): upgrade pytest-docker (#9765)

---
 metadata-ingestion/setup.py                   |  2 +-
 .../tests/test_helpers/docker_helpers.py      | 20 +++++++++++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
index f8d5199733..d4e2ada1fc 100644
--- a/metadata-ingestion/setup.py
+++ b/metadata-ingestion/setup.py
@@ -468,7 +468,7 @@ base_dev_requirements = {
     pytest_dep,
     "pytest-asyncio>=0.16.0",
     "pytest-cov>=2.8.1",
-    "pytest-docker>=1.0.1",
+    "pytest-docker>=1.1.0",
     deepdiff_dep,
     "requests-mock",
     "freezegun",
diff --git a/metadata-ingestion/tests/test_helpers/docker_helpers.py b/metadata-ingestion/tests/test_helpers/docker_helpers.py
index 2eb6106819..bacb8d80b9 100644
--- a/metadata-ingestion/tests/test_helpers/docker_helpers.py
+++ b/metadata-ingestion/tests/test_helpers/docker_helpers.py
@@ -2,7 +2,7 @@ import contextlib
 import logging
 import os
 import subprocess
-from typing import Callable, Optional, Union
+from typing import Callable, Iterator, List, Optional, Union
 
 import pytest
 import pytest_docker.plugin
@@ -37,9 +37,11 @@ def wait_for_port(
         docker_services.wait_until_responsive(
             timeout=timeout,
             pause=pause,
-            check=checker
-            if checker
-            else lambda: is_responsive(container_name, container_port, hostname),
+            check=(
+                checker
+                if checker
+                else lambda: is_responsive(container_name, container_port, hostname)
+            ),
         )
         logger.info(f"Container {container_name} is ready!")
     finally:
@@ -62,14 +64,16 @@ def docker_compose_runner(
 ):
     @contextlib.contextmanager
     def run(
-        compose_file_path: Union[str, list], key: str, cleanup: bool = True
-    ) -> pytest_docker.plugin.Services:
+        compose_file_path: Union[str, List[str]], key: str, cleanup: bool = True
+    ) -> Iterator[pytest_docker.plugin.Services]:
         with pytest_docker.plugin.get_docker_services(
             docker_compose_command=docker_compose_command,
-            docker_compose_file=compose_file_path,
+            # We can remove the type ignore once this is merged:
+            # https://github.com/avast/pytest-docker/pull/108
+            docker_compose_file=compose_file_path,  # type: ignore
             docker_compose_project_name=f"{docker_compose_project_name}-{key}",
             docker_setup=docker_setup,
-            docker_cleanup=docker_cleanup if cleanup else False,
+            docker_cleanup=docker_cleanup if cleanup else [],
         ) as docker_services:
             yield docker_services