Skip to content

Commit

Permalink
Adds validation to image_spec for list of strings (flyteorg#2655)
Browse files Browse the repository at this point in the history
Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>
Signed-off-by: mao3267 <chenvincent610@gmail.com>
  • Loading branch information
thomasjpfan authored and mao3267 committed Aug 9, 2024
1 parent 8d229d9 commit 59b5932
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 0 deletions.
17 changes: 17 additions & 0 deletions flytekit/image_spec/image_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,23 @@ def __post_init__(self):
if self.registry:
self.registry = self.registry.lower()

parameters_str_list = [
"packages",
"conda_channels",
"conda_packages",
"apt_packages",
"pip_extra_index_url",
"entrypoint",
"commands",
]
for parameter in parameters_str_list:
attr = getattr(self, parameter)
parameter_is_None = attr is None
parameter_is_list_string = isinstance(attr, list) and all(isinstance(v, str) for v in attr)
if not (parameter_is_None or parameter_is_list_string):
error_msg = f"{parameter} must be a list of strings or None"
raise ValueError(error_msg)

def image_name(self) -> str:
"""Full image name with tag."""
image_name = self._image_name()
Expand Down
16 changes: 16 additions & 0 deletions tests/flytekit/unit/core/image_spec/test_image_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,3 +138,19 @@ def test_no_build_during_execution():
ImageBuildEngine.build(spec)

ImageBuildEngine._build_image.assert_not_called()


@pytest.mark.parametrize(
"parameter_name", [
"packages", "conda_channels", "conda_packages",
"apt_packages", "pip_extra_index_url", "entrypoint", "commands"
]
)
@pytest.mark.parametrize("value", ["requirements.txt", [1, 2, 3]])
def test_image_spec_validation_string_list(parameter_name, value):
msg = f"{parameter_name} must be a list of strings or None"

input_params = {parameter_name: value}

with pytest.raises(ValueError, match=msg):
ImageSpec(**input_params)

0 comments on commit 59b5932

Please sign in to comment.