Function mock_data_source_state

tests/helpers/airflow_tests/test_airflow_wrapper.py:112–162

()

@dlt.source(section="mock_data_source_state")
def mock_data_source_state():
    @dlt.resource(selected=True)
    def _r_init():
        assert is_running_in_airflow_task() is True, "Must be running in Airflow task"
        # will switch to spawn on Airflow
        assert get_default_start_method("fork") == "spawn"

        dlt.current.source_state()["counter"] = 1
        dlt.current.source_state()["end_counter"] = 1
        yield ["-", "x", "!"]

    @dlt.resource(selected=False)
    def _r1():
        dlt.current.source_state()["counter"] += 1
        dlt.current.resource_state()["counter"] = 1
        yield from ["a", "b", "c"]

    @dlt.transformer(data_from=_r1, selected=True)
    def _t1(items, suffix):
        dlt.current.source_state()["counter"] += 1
        dlt.current.resource_state("_r1")["counter"] += 1
        dlt.current.resource_state()["counter"] = 1
        yield list(map(lambda i: i + "_" + suffix, items))

    @dlt.transformer(data_from=_r1)
    def _t2(items, mul):
        dlt.current.source_state()["counter"] += 1
        dlt.current.resource_state("_r1")["counter"] += 1
        dlt.current.resource_state()["counter"] = 1
        yield items * mul

    @dlt.transformer(data_from=_r1)
    def _t3(items, mul):
        dlt.current.source_state()["counter"] += 1
        dlt.current.resource_state("_r1")["counter"] += 1
        dlt.current.resource_state()["counter"] = 1
        for item in items:
            yield item.upper() * mul

    # add something to init
    @dlt.transformer(data_from=_r_init)
    def _t_init_post(items):
        for item in items:
            yield item * 2

    @dlt.resource
    def _r_isolee():
        dlt.current.source_state()["end_counter"] += 1
        yield from ["AX", "CV", "ED"]

    return _r_init, _t_init_post, _r1, _t1("POST"), _t2(3), _t3(2), _r_isolee


def test_regular_run() -> None: