[nit] Add unit test

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -1,11 +1,16 @@
 import os
 import pathlib
 
+import mock
 import pytest
 from click.testing import CliRunner
 
 from flytekit.clis.sdk_in_container import pyflyte
-from flytekit.clis.sdk_in_container.run import get_entities_in_file
+from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE
+from flytekit.clis.sdk_in_container.helpers import FLYTE_REMOTE_INSTANCE_KEY
+from flytekit.clis.sdk_in_container.run import REMOTE_FLAG_KEY, RUN_LEVEL_PARAMS_KEY, get_entities_in_file, run_command
+from flytekit.configuration import Image, ImageConfig
+from flytekit.core.task import task
 
 WORKFLOW_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "workflow.py")
 DIR_NAME = os.path.dirname(os.path.realpath(__file__))
@@ -170,5 +175,69 @@ def test_list_default_arguments(wf_path):
         ],
         catch_exceptions=False,
     )
-    print(result.stdout)
     assert result.exit_code == 0
+
+
+# default case, what comes from click if no image is specified, the click param is configured to use the default.
+ic_result_1 = ImageConfig(
+    default_image=Image(name="default", fqn="ghcr.io/flyteorg/mydefault", tag="py3.9-latest"),
+    images=[Image(name="default", fqn="ghcr.io/flyteorg/mydefault", tag="py3.9-latest")],
+)
+# test that command line args are merged with the file
+ic_result_2 = ImageConfig(
+    default_image=None,
+    images=[
+        Image(name="asdf", fqn="ghcr.io/asdf/asdf", tag="latest"),
+        Image(name="xyz", fqn="docker.io/xyz", tag="latest"),
+        Image(name="abc", fqn="docker.io/abc", tag=None),
+    ],
+)
+# test that command line args override the file
+ic_result_3 = ImageConfig(
+    default_image=None,
+    images=[Image(name="xyz", fqn="ghcr.io/asdf/asdf", tag="latest"), Image(name="abc", fqn="docker.io/abc", tag=None)],
+)
+
+
+@pytest.mark.parametrize(
+    "image_string, leaf_configuration_file_name, final_image_config",
+    [
+        ("ghcr.io/flyteorg/mydefault:py3.9-latest", "no_images.yaml", ic_result_1),
+        ("asdf=ghcr.io/asdf/asdf:latest", "sample.yaml", ic_result_2),
+        ("xyz=ghcr.io/asdf/asdf:latest", "sample.yaml", ic_result_3),
+    ],
+)
+def test_pyflyte_run_run(image_string, leaf_configuration_file_name, final_image_config):
+    @task
+    def a():
+        ...
+
+    mock_click_ctx = mock.MagicMock()
+    mock_remote = mock.MagicMock()
+    image_tuple = (image_string,)
+    image_config = ImageConfig.validate_image(None, "", image_tuple)
+
+    run_level_params = {
+        "project": "p",
+        "domain": "d",
+        "image_config": image_config,
+    }
+
+    pp = pathlib.Path.joinpath(
+        pathlib.Path(__file__).parent.parent.parent, "configuration/configs/", leaf_configuration_file_name
+    )
+
+    obj = {
+        RUN_LEVEL_PARAMS_KEY: run_level_params,
+        REMOTE_FLAG_KEY: True,
+        FLYTE_REMOTE_INSTANCE_KEY: mock_remote,
+        CTX_CONFIG_FILE: str(pp),
+    }
+    mock_click_ctx.obj = obj
+
+    def check_image(*args, **kwargs):
+        assert kwargs["image_config"] == final_image_config
+
+    mock_remote.register_script.side_effect = check_image
+
+    run_command(mock_click_ctx, a)()