You’re also welcome to leave details here

Here’s the gist of it: 1. I have a

HdfsSparkDataFrameResult

that appears to work as expected when testing its individual

write

,

read

, and

exists

methods in unit tests. 2. I’m now trying to test that Prefect is properly using that result when re-running that task. 3. To do that, I’ve decorated

select_one

with

result=HdfsSparkDataFrameResult(...)

, which – if the result works as expected – I imagine would first check whether the result exists, and if it does, reads the DataFrame from HDFS. 4. I’m not sure what assertions to make – either on the

flow

, or the flow state

fs

, or on the task result

result

, or on the re-run flow state

fs2

– to assert that Prefect is just running the task once (on the first

flow.run()

) and reading from cache on the second

flow.run()

.

from internalpackagename.hdfs import HDFSPath
from loguru import logger
from prefect.core.flow import Flow
from prefect.engine.flow_runner import FlowRunner
from prefect.engine.state import State
from prefect.utilities.tasks import task
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.session import SparkSession

from prefectx.results import HdfsSparkDataFrameResult


class TestHdfsSparkDataFrameResult:

    def test_result_on_task(self, tmp_hdfs_dir: HDFSPath, spark: SparkSession) -> None:

        # Create a `prefect` task that outputs a PySpark DataFrame.
        @task(
            result=HdfsSparkDataFrameResult(
                location=tmp_hdfs_dir,
                write_kwargs=dict(mode="overwrite")
            )
        )
        def select_one() -> DataFrame:
            """Create and return a placeholder DataFrame."""
            df: DataFrame = spark.sql(sqlQuery="SELECT 1")
            return df

        # Run our flow once, knowing that the DataFrame won't be cached.
        with Flow(name="test") as flow:
            result = select_one()
        fr: FlowRunner = FlowRunner(flow=flow)
        fs: State = fr.run()

        # First, make sure the task/flow was successful.
        assert fs.is_successful()

        # Maybe re-run the flow so that the task gets cached?
        fs2: State = fr.run()
        assert fs2.is_successful()
        # Maybe test that the flow was cached? No, right? The task should be cached?
        assert fs2.is_cached()

I think I’d want to test that the

select_one

task’s state on

fs2

would be cached (instead of successful), but I’m not 100% sure.

Quick edit: to be safe, I think we need to pass

checkpoint=True

to the

@task

decorator. In my testing environment I already have

PREFECT__FLOWS__CHECKPOINTING=true

so that’s good 👍

Checkpointing is about Flow restarts and continuing from where work left off while caching described like “don’t run this task for flow runs in the next 24 hours”

I think this test is clear and might give you an idea for what you want to do https://github.com/PrefectHQ/prefect/blob/master/tests/engine/test_task_runner.py#L1217-L1230