fix(ingest): cleanup large images in CI (#9153)

metadata-ingestion/tests/integration/sql_server/test_sql_server.py

@@ -6,7 +6,7 @@
 
 from tests.test_helpers import mce_helpers
 from tests.test_helpers.click_helpers import run_datahub_cmd
-from tests.test_helpers.docker_helpers import wait_for_port
+from tests.test_helpers.docker_helpers import cleanup_image, wait_for_port
 
 
 @pytest.fixture(scope="module")
@@ -29,6 +29,9 @@ def mssql_runner(docker_compose_runner, pytestconfig):
         assert ret.returncode == 0
         yield docker_services
 
+    # The image is pretty large, so we remove it after the test.
+    cleanup_image("mcr.microsoft.com/mssql/server")
+
 
 SOURCE_FILES_PATH = "./tests/integration/sql_server/source_files"
 config_file = os.listdir(SOURCE_FILES_PATH)

metadata-ingestion/tests/integration/vertica/test_vertica.py

@@ -6,7 +6,7 @@
 
 from tests.test_helpers import mce_helpers
 from tests.test_helpers.click_helpers import run_datahub_cmd
-from tests.test_helpers.docker_helpers import wait_for_port
+from tests.test_helpers.docker_helpers import cleanup_image, wait_for_port
 
 FROZEN_TIME = "2020-04-14 07:00:00"
 
@@ -49,6 +49,9 @@ def vertica_runner(docker_compose_runner, test_resources_dir):
 
         yield docker_services
 
+    # The image is pretty large, so we remove it after the test.
+    cleanup_image("vertica/vertica-ce")
+
 
 @freeze_time(FROZEN_TIME)
 @pytest.mark.integration

metadata-ingestion/tests/test_helpers/docker_helpers.py

@@ -1,5 +1,6 @@
 import contextlib
 import logging
+import os
 import subprocess
 from typing import Callable, Optional, Union
 
@@ -78,6 +79,10 @@ def run(
 def cleanup_image(image_name: str) -> None:
     assert ":" not in image_name, "image_name should not contain a tag"
 
+    if not os.environ.get("CI"):
+        logger.debug("Not cleaning up images to speed up local development")
+        return
+
     images_proc = subprocess.run(
         f"docker image ls --filter 'reference={image_name}*' -q",
         shell=True,