diff --git a/tests/flytekit/unit/core/test_flyte_file.py b/tests/flytekit/unit/core/test_flyte_file.py index 7abb3365aa..a12c414f35 100644 --- a/tests/flytekit/unit/core/test_flyte_file.py +++ b/tests/flytekit/unit/core/test_flyte_file.py @@ -54,21 +54,25 @@ def can_import(module_name) -> bool: def test_file_type_in_workflow_with_bad_format(): - @task - def t1() -> FlyteFile[typing.TypeVar("txt")]: - fname = "/tmp/flytekit_test" - with open(fname, "w") as fh: - fh.write("Hello World\n") - return fname + fd, path = tempfile.mkstemp() + try: - @workflow - def my_wf() -> FlyteFile[typing.TypeVar("txt")]: - f = t1() - return f + @task + def t1() -> FlyteFile[typing.TypeVar("txt")]: + with os.fdopen(fd, "w") as f: + f.write("Hello World\n") + return path - res = my_wf() - with open(res, "r") as fh: - assert fh.read() == "Hello World\n" + @workflow + def my_wf() -> FlyteFile[typing.TypeVar("txt")]: + f = t1() + return f + + res = my_wf() + with open(res, "r") as fh: + assert fh.read() == "Hello World\n" + finally: + os.remove(path) def test_matching_file_types_in_workflow(local_dummy_txt_file):