import pytest

from dagster import (
    DagsterInvariantViolationError,
    ModeDefinition,
    execute_pipeline,
    pipeline,
    resource,
    solid,
)


@resource
def add_one_resource(_):
    def add_one(num):
        return num + 1

    return add_one


@resource
def add_two_resource(_):
    def add_two(num):
        return num + 2

    return add_two


@solid(required_resource_keys={"adder"})
def solid_that_uses_adder_resource(context, number):
    return context.resources.adder(number)


@pipeline(
    mode_defs=[
        ModeDefinition(name="add_one", resource_defs={"adder": add_one_resource}),
        ModeDefinition(name="add_two", resource_defs={"adder": add_two_resource}),
    ]
)
def pipeline_with_mode():
    return solid_that_uses_adder_resource()


def test_execute_pipeline_with_mode():
    pipeline_result = execute_pipeline(
        pipeline_with_mode,
        run_config={
            "solids": {"solid_that_uses_adder_resource": {"inputs": {"number": {"value": 4}}}}
        },
        mode="add_one",
    )
    assert pipeline_result.success
    assert pipeline_result.result_for_solid("solid_that_uses_adder_resource").output_value() == 5

    pipeline_result = execute_pipeline(
        pipeline_with_mode,
        run_config={
            "solids": {"solid_that_uses_adder_resource": {"inputs": {"number": {"value": 4}}}}
        },
        mode="add_two",
    )
    assert pipeline_result.success
    assert pipeline_result.result_for_solid("solid_that_uses_adder_resource").output_value() == 6


def test_execute_pipeline_with_non_existant_mode():
    with pytest.raises(DagsterInvariantViolationError):
        execute_pipeline(
            pipeline_with_mode,
            mode="BAD",
            run_config={
                "solids": {"solid_that_uses_adder_resource": {"inputs": {"number": {"value": 4}}}}
            },
        )


@solid
def solid_that_gets_tags(context):
    return context.pipeline_run.tags


@pipeline(
    mode_defs=[ModeDefinition(name="tags"),], tags={"tag_key": "tag_value"},
)
def pipeline_with_one_mode_and_tags():
    return solid_that_gets_tags()


def test_execute_pipeline_with_mode_and_tags():
    pipeline_result = execute_pipeline(pipeline_with_one_mode_and_tags)
    assert pipeline_result.success
    assert pipeline_result.result_for_solid("solid_that_gets_tags").output_value() == {
        "tag_key": "tag_value"
    }


@pipeline(
    mode_defs=[
        ModeDefinition(name="tags_1"),
        ModeDefinition(name="tags_2"),
        ModeDefinition(name="tags_3"),
    ],
    tags={"pipeline_tag_key": "pipeline_tag_value"},
)
def pipeline_with_multi_mode_and_tags():
    return solid_that_gets_tags()


def test_execute_pipeline_with_multi_mode_and_pipeline_def_tags():
    pipeline_result = execute_pipeline(pipeline_with_multi_mode_and_tags, mode="tags_1")
    assert pipeline_result.success
    assert pipeline_result.result_for_solid("solid_that_gets_tags").output_value() == {
        "pipeline_tag_key": "pipeline_tag_value"
    }


def test_execute_pipeline_with_multi_mode_and_pipeline_def_tags_and_execute_tags():
    pipeline_result = execute_pipeline(
        pipeline_with_multi_mode_and_tags, mode="tags_1", tags={"run_tag_key": "run_tag_value"},
    )
    assert pipeline_result.success
    assert pipeline_result.result_for_solid("solid_that_gets_tags").output_value() == {
        "pipeline_tag_key": "pipeline_tag_value",
        "run_tag_key": "run_tag_value",
    }
