Adds test security contexts and fixes test cases

Signed-off-by: Geert Pingen <[email protected]>
flyteorg · Jul 8, 2023 · 2154a2d · 2154a2d
1 parent 835719e
commit 2154a2d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git a/Makefile b/Makefile
@@ -57,8 +57,9 @@ unit_test_codecov:
 unit_test:
 	# Skip tensorflow tests and run them with the necessary env var set so that a working (albeit slower)
 	# library is used to serialize/deserialize protobufs is used.
+	# Can use pytest --lf to only rerun previously failed tests.
 	pytest -m "not sandbox_test" tests/flytekit/unit/ --ignore=tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS} && \
-		PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python3 pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS}
+		PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python pytest tests/flytekit/unit/extras/tensorflow ${CODECOV_OPTS}
 
 doc-requirements.txt: export CUSTOM_COMPILE_COMMAND := make doc-requirements.txt
 doc-requirements.txt: doc-requirements.in install-piptools

diff --git a/flytekit/models/security.py b/flytekit/models/security.py
@@ -67,12 +67,11 @@ def from_flyte_idl(cls, pb2_object: _sec.Secret.MountFile):
             )
 
     group: str
-    env_var: Optional[MountEnvVar] = None
-    file: Optional[MountFile] = None
-
     key: Optional[str] = None
     group_version: Optional[str] = None
     mount_requirement: MountType = MountType.ANY
+    env_var: Optional[MountEnvVar] = field(default_factory=lambda: None)
+    file: Optional[MountFile] = field(default_factory=lambda: None)
 
     def __post_init__(self):
         if self.group is None:
@@ -84,8 +83,8 @@ def to_flyte_idl(self) -> _sec.Secret:
             group_version=self.group_version,
             key=self.key,
             mount_requirement=self.mount_requirement.value,
-            env_var=self.env_var.to_flyte_idl() if self.env_var else None,
-            file=self.file.to_flyte_idl() if self.file else None,
+            env_var=_sec.Secret.MountEnvVar(name=self.env_var.name) if self.env_var else None,
+            file=_sec.Secret.MountFile(path=self.file.path) if self.file else None,
         )
 
     @classmethod

diff --git a/tests/flytekit/common/parameterizers.py b/tests/flytekit/common/parameterizers.py
@@ -238,7 +238,9 @@
     None,
     security.Secret(group="x", key="g"),
     security.Secret(group="x", key="y", mount_requirement=security.Secret.MountType.ANY),
-    security.Secret(group="x", key="y", mount_target=security.Secret.MountEnvVar(name="z"), mount_requirement=security.Secret.MountType.ANY),
+    security.Secret(group="x", key="y", env_var=security.Secret.MountEnvVar(name="z")),
+    security.Secret(group="x", key="y", file=security.Secret.MountFile(path="/z")),
+    security.Secret(group="x", key="y", group_version="1", mount_requirement=security.Secret.MountType.ENV_VAR),
     security.Secret(group="x", key="y", group_version="1", mount_requirement=security.Secret.MountType.FILE),
 ]