From f5a72f7d02d457bdc59439af03a031e583d67c19 Mon Sep 17 00:00:00 2001 From: JiangJiaWei1103 Date: Wed, 29 Jan 2025 15:24:42 +0800 Subject: [PATCH] test: Test image attr in task and dynamic decorators Signed-off-by: JiangJiaWei1103 --- flytekit/core/task.py | 2 +- tests/flytekit/unit/core/test_task.py | 50 +++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 86cac2743a..1453759e9a 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -380,7 +380,7 @@ def wrapper(fn: Callable[P, Any]) -> PythonFunctionTask[T]: _container_image = None if image is not None and container_image is not None: raise ValueError( - "Cannot specify both image and container_image." + "Cannot specify both image and container_image. " "Please use image because container_image is deprecated and will be removed in the future." ) elif container_image is not None: diff --git a/tests/flytekit/unit/core/test_task.py b/tests/flytekit/unit/core/test_task.py index b02f7cb08b..592db7eee5 100644 --- a/tests/flytekit/unit/core/test_task.py +++ b/tests/flytekit/unit/core/test_task.py @@ -1,11 +1,16 @@ +import os import pytest +from flytekit import task, dynamic from flytekit.core.task import decorate_function from flytekit.core.utils import str2bool from flytekit.interactive import vscode from flytekit.interactive.constants import FLYTE_ENABLE_VSCODE_KEY +IMAGE = os.environ.get("FLYTEKIT_IMAGE", "localhost:30000/flytekit:dev") + + def test_decorate_python_task(monkeypatch: pytest.MonkeyPatch): def t1(a: int, b: int) -> int: return a + b @@ -13,3 +18,48 @@ def t1(a: int, b: int) -> int: assert decorate_function(t1) is t1 monkeypatch.setenv(FLYTE_ENABLE_VSCODE_KEY, str2bool("True")) assert isinstance(decorate_function(t1), vscode) + + +def test_image(): + # Define expected warning and error messages + WARN_MSG = "container_image is deprecated and will be removed in the future. Please use image instead." + ERR_MSG = ( + "Cannot specify both image and container_image. " + "Please use image because container_image is deprecated and will be removed in the future." + ) + + # Plain tasks + @task(image=IMAGE) + def t1() -> str: + return "Use image in @task." + assert t1._container_image == IMAGE, f"_container_image of t1 should match the user-specified {IMAGE}" + + with pytest.warns(DeprecationWarning, match=WARN_MSG): + @task(container_image=IMAGE) + def t2() -> str: + return "Use container_image in @task." + assert t2._container_image == IMAGE, f"_container_image of t2 should match the user-specified {IMAGE}" + + with pytest.raises(ValueError, match=ERR_MSG): + @task(image=IMAGE, container_image=IMAGE) + def t3() -> str: + return "Use both image and container_image in @task." + + # Dynamic workflow tasks + @dynamic(image=IMAGE) + def dy_t1(i: int) -> str: + return "Use image in @dynamic." + assert dy_t1._container_image == IMAGE, f"_container_image of dy_t1 should match the user-specified {IMAGE}" + + with pytest.warns(DeprecationWarning, match=WARN_MSG): + @dynamic(container_image=IMAGE) + def dy_t2(i: int) -> str: + return "Use container_image in @dynamic." + assert dy_t2._container_image == IMAGE, f"_container_image of dy_t2 should match the user-specified {IMAGE}" + + with pytest.raises(ValueError, match=ERR_MSG): + @dynamic(image=IMAGE, container_image=IMAGE) + def dy_t3(i: int) -> str: + return "Use both image and container_image in @dynamic." + + # TODO: Test eager