Follow FLYTE_PUSH_IMAGE_SPEC in default image builder (#2682)

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>

flytekit/image_spec/default_builder.py

@@ -1,12 +1,13 @@
 import json
+import os
 import re
 import shutil
-import subprocess
 import sys
 import tempfile
 import warnings
 from pathlib import Path
 from string import Template
+from subprocess import run
 from typing import ClassVar
 
 import click
@@ -251,7 +252,10 @@ class DefaultImageBuilder(ImageSpecBuilder):
     }
 
     def build_image(self, image_spec: ImageSpec) -> str:
-        return self._build_image(image_spec)
+        return self._build_image(
+            image_spec,
+            push=os.getenv("FLYTE_PUSH_IMAGE_SPEC", "True").lower() in ("true", "1"),
+        )
 
     def _build_image(self, image_spec: ImageSpec, *, push: bool = True) -> str:
         # For testing, set `push=False`` to just build the image locally and not push to
@@ -285,4 +289,4 @@ def _build_image(self, image_spec: ImageSpec, *, push: bool = True) -> str:
 
             concat_command = " ".join(command)
             click.secho(f"Run command: {concat_command} ", fg="blue")
-            subprocess.run(command, check=True)
+            run(command, check=True)

tests/flytekit/unit/core/image_spec/test_default_builder.py

@@ -1,4 +1,5 @@
 import os
+from unittest.mock import patch, Mock
 
 import pytest
 
@@ -181,3 +182,23 @@ def test_build(tmp_path):
     builder = DefaultImageBuilder()
 
     builder.build_image(image_spec)
+
+
+@pytest.mark.parametrize("push_image_spec", ["0", "1"])
+def test_should_push_env(monkeypatch, push_image_spec):
+    image_spec = ImageSpec(name="my_flytekit", python_version="3.12", registry="localhost:30000")
+    monkeypatch.setenv("FLYTE_PUSH_IMAGE_SPEC", push_image_spec)
+
+    run_mock = Mock()
+    monkeypatch.setattr("flytekit.image_spec.default_builder.run", run_mock)
+
+    builder = DefaultImageBuilder()
+    builder.build_image(image_spec)
+
+    run_mock.assert_called_once()
+    call_args = run_mock.call_args.args
+
+    if push_image_spec == "0":
+        assert "--push" not in call_args[0]
+    else:
+        assert "--push" in call_args[0]