diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 32abcc2dd2..89bb8bd1b3 100644 --- a/flytekit/image_spec/default_builder.py +++ b/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) diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 5d839e0f39..e61a3cb7c8 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/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]