@ancient-wolf-19325 I am writing tests for tasks. One of the tasks returns a CSVFile that writes to

fl.current_context().working_directory

/ "my_file.csv". In my tests I actually don't want to save it permanently, but to

tmp_path

so that I don't clutter my storage. How do I do this?

It seems likely that you can achieve this by using a temporary directory for your tests. You can use Python's

tempfile

module to create a temporary directory and set it as the working directory for your task during testing. This way, any files created by the task will be stored in the temporary directory, which will be automatically cleaned up after the test. Here's a basic example of how you might set this up in your test:

import tempfile
from flytekit import task, current_context
from flytekit.types.file import CSVFile

@task
def create_csv_file() -> CSVFile:
    ctx = current_context()
    wd = ctx.working_directory
    file_path = f"{wd}/my_file.csv"
    with open(file_path, "w") as f:
        f.write("column1,column2\nvalue1,value2")
    return CSVFile(file_path)

def test_create_csv_file():
    with tempfile.TemporaryDirectory() as tmp_dir:
        # Set the working directory to the temporary directory
        ctx = current_context().with_working_directory(tmp_dir)
        with ctx:
            csv_file = create_csv_file()
            assert csv_file.path.startswith(tmp_dir)
            # Further assertions can be added here

This example uses a context manager to temporarily set the working directory to a temporary directory created by

tempfile.TemporaryDirectory()

. The file will be written to this temporary directory, and it will be cleaned up automatically when the context manager exits. Would you like more details on this approach or any other aspect of testing with Flyte? Sources: - https://github.com/flyteorg/flytekit/blob/master/tests/flytekit/unit/core/test_generic_idl_protobuf_struct.py 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.

@ancient-wolf-19325

with_working_directory

does not exist

It seems likely that the

with_working_directory

method is not available in the Flytekit API. Instead, you can manually set the working directory in your test by modifying the context's

working_directory

attribute directly. Here's how you can do it:

import tempfile
from flytekit import task, current_context
from flytekit.types.file import CSVFile

@task
def create_csv_file() -> CSVFile:
    ctx = current_context()
    wd = ctx.working_directory
    file_path = f"{wd}/my_file.csv"
    with open(file_path, "w") as f:
        f.write("column1,column2\nvalue1,value2")
    return CSVFile(file_path)

def test_create_csv_file():
    with tempfile.TemporaryDirectory() as tmp_dir:
        # Modify the current context to use the temporary directory as the working directory
        ctx = current_context()
        ctx.working_directory = tmp_dir
        with ctx:
            csv_file = create_csv_file()
            assert csv_file.path.startswith(tmp_dir)
            # Further assertions can be added here

This approach manually sets the

working_directory

attribute of the current context to the temporary directory created by

tempfile.TemporaryDirectory()

. This way, the file will be written to the temporary directory, and it will be cleaned up automatically when the context manager exits. Would you like more details on this approach or any other aspect of testing with Flyte? Sources: - https://github.com/flyteorg/flyte/issues/4638 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.

@ancient-wolf-19325

working_directory

can not be set as it is a property

It seems likely that directly setting the

working_directory

property is not possible because it is a read-only attribute. Instead, you can use the

FlyteContextManager

to create a new context with a modified working directory. Here's how you can achieve this:

import tempfile
from flytekit import task, FlyteContextManager
from flytekit.types.file import CSVFile

@task
def create_csv_file() -> CSVFile:
    ctx = FlyteContextManager.current_context()
    wd = ctx.working_directory
    file_path = f"{wd}/my_file.csv"
    with open(file_path, "w") as f:
        f.write("column1,column2\nvalue1,value2")
    return CSVFile(file_path)

def test_create_csv_file():
    with tempfile.TemporaryDirectory() as tmp_dir:
        # Create a new context with the temporary directory as the working directory
        ctx = FlyteContextManager.current_context()
        new_ctx = ctx.with_execution_state(ctx.execution_state.with_params(working_dir=tmp_dir))
        with FlyteContextManager.with_context(new_ctx):
            csv_file = create_csv_file()
            assert csv_file.path.startswith(tmp_dir)
            # Further assertions can be added here

This approach uses

FlyteContextManager.with_context

to temporarily switch to a new context where the

working_directory

is set to the temporary directory. This ensures that any files created during the test are stored in the temporary directory, which is cleaned up automatically. Would you like more details on this approach or any other aspect of testing with Flyte? Sources: - https://github.com/flyteorg/flyte/issues/4638 - https://github.com/flyteorg/flytekit/blob/master/tests/flytekit/unit/types/directory/test_listdir.py 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.

@ancient-wolf-19325 How do I clear the folders and files created in

fl.current_context().working_directory

?

It seems likely that you can clear the folders and files created in

fl.current_context().working_directory

by using Python's

shutil

module to remove the directory contents. Here's a basic example of how you might do this:

import shutil
from flytekit import FlyteContextManager

def clear_working_directory():
    ctx = FlyteContextManager.current_context()
    wd = ctx.working_directory
    # Remove all contents of the working directory
    shutil.rmtree(wd, ignore_errors=True)
    # Recreate the working directory
    os.makedirs(wd, exist_ok=True)

# Example usage
clear_working_directory()

This function uses

shutil.rmtree

to remove all contents of the working directory, and then recreates the directory to ensure it exists for future operations. The

ignore_errors=True

parameter ensures that any errors during the removal process are ignored, which can be useful if some files are in use or if the directory is already empty. Would you like more details on this approach or any other aspect of managing the working directory in Flyte? Sources: - https://github.com/flyteorg/flytekit/blob/master/tests/flytekit/unit/core/test_flyte_directory.py 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.