def test_get_dsk_attributes_wrong_color_dtype():
    dsk = {char: None for char in ["A"]}

    df_tasks = pd.DataFrame(
        [{"key": "A", "action": "compute", "color": datetime.datetime.today()},]
    )

    with pytest.raises(ValueError):
        _get_dsk_attributes(
            dsk, df_tasks, color_col="color",
        )
def test_get_dsk_attributes_labels():
    dsk = {
        "A": None,
    }

    df_tasks = pd.DataFrame([{"key": "A", "action": "compute", "label": "label"},])

    dsk_attributes = _get_dsk_attributes(dsk, df_tasks, label_col="label")
    dsk_attributes_test = {
        "func": {"A": {"label": "label"},},
        "data": {},
    }
    assert dsk_attributes == dsk_attributes_test
def test_get_dsk_attributes_category_color():
    dsk = {char: None for char in ["A", "B"]}

    df_tasks = pd.DataFrame(
        [
            {"key": "A", "action": "compute", "color": "Z"},
            {"key": "B", "action": "transfer", "color": "Y"},
        ]
    )

    dsk_attributes = _get_dsk_attributes(dsk, df_tasks, color_col="color")
    dsk_attributes_test = {
        "func": {"A": {"style": "filled", "fillcolor": "#10AEFD"}},
        "data": {"B": {"style": "filled", "fillcolor": "#DB9758"}},
    }
    assert dsk_attributes == dsk_attributes_test
def test_get_dsk_attributes_float_color():
    dsk = {char: None for char in ["A", "B"]}

    df_tasks = pd.DataFrame(
        [
            {"key": "A", "action": "compute", "color": 20},
            {"key": "B", "action": "transfer", "color": 5.5},
        ]
    )

    dsk_attributes = _get_dsk_attributes(dsk, df_tasks, color_col="color")
    dsk_attributes_test = {
        "func": {"A": {"style": "filled", "fillcolor": "gray0", "fontcolor": "white"}},
        "data": {"B": {"style": "filled", "fillcolor": "gray73"}},
    }

    assert dsk_attributes == dsk_attributes_test
def test_get_dsk_attributes_na_categorical_color():
    dsk = {char: None for char in ["A", "B", "C"]}

    df_tasks = pd.DataFrame(
        [
            {"key": "A", "action": "compute"},
            {"key": "B", "action": "compute", "color": "color"},
            {"key": "C", "action": "compute"},
        ]
    )

    dsk_attributes = _get_dsk_attributes(dsk, df_tasks, color_col="color",)
    dsk_attributes_test = {
        "func": {"B": {"fillcolor": "#10AEFD", "style": "filled"},},
        "data": {},
    }

    assert dsk_attributes == dsk_attributes_test
def test_get_dsk_attributes_na_label():
    dsk = {char: None for char in ["A", "B", "C"]}

    df_tasks = pd.DataFrame(
        [
            {"key": "A", "action": "compute"},
            {"key": "B", "action": "compute", "label": "label"},
            {"key": "C", "action": "compute"},
        ]
    )

    dsk_attributes = _get_dsk_attributes(dsk, df_tasks, label_col="label",)
    dsk_attributes_test = {
        "func": {"B": {"label": "label"},},
        "data": {},
    }

    assert dsk_attributes == dsk_attributes_test
def test_get_dsk_attributes_progress_color():
    dsk = {char: None for char in ["A", "B"]}

    df_tasks = pd.DataFrame(
        [
            {"key": "A", "action": "compute", "start_delta": 1, "stop_delta": 3},
            {"key": "B", "action": "transfer", "start_delta": 2, "stop_delta": 5},
            {"key": "C", "action": "compute", "start_delta": 6, "stop_delta": 7},
        ]
    )

    dsk_attributes = _get_dsk_attributes(
        dsk, df_tasks, color_col="progress", current_time=4,
    )
    dsk_attributes_test = {
        "func": {"A": {"style": "filled", "fillcolor": "blue"}},
        "data": {"B": {"style": "filled", "fillcolor": "red"}},
    }
    assert dsk_attributes == dsk_attributes_test