from assets import asset_multi_1, asset_multi_2, asset_single_1, asset_single_2
from dagster_polars import PolarsParquetIOManager

from dagster import materialize

resources = {
    "polars_parquet_io_manager": PolarsParquetIOManager(base_dir="/opt/dagster/storage")
}


def test_single():
    result = materialize(
        [asset_single_1, asset_single_2],
        partition_key="2024-10-02",
        resources=resources,
    )
    assert result.success
    ic(result.asset_value)


def test_multi():
    # result = materialize([
    #     asset_multi_1
    # ], partition_key="2024-10-01|plato", resources=resources
    # )
    # assert result.success
    # ic(result.asset_value)
    #
    #
    # result = materialize([
    #     asset_multi_1
    # ], partition_key="2024-10-02|plato", resources=resources
    # )
    # assert result.success
    # ic(result.asset_value)

    result = materialize(
        [asset_multi_1, asset_multi_2],
        partition_key="2024-10-02|plato",
        resources=resources,
    )
    assert result.success
    ic(result.asset_value)


if __name__ == "__main__":
    # test_single()
    test_multi()