from airflow_pydantic import Dag


class TestRender:
    def test_render_operator_ssh(self, ssh_operator):
        imports, globals_, task = ssh_operator.render()
        assert imports == [
            "from airflow.providers.ssh.operators.ssh import SSHOperator",
            "from airflow.providers.ssh.hooks.ssh import SSHHook",
        ]
        assert globals_ == []
        assert (
            task
            == "SSHOperator(do_xcom_push=True, ssh_hook=SSHHook(remote_host='test', username='test'), ssh_conn_id='test', command='test', cmd_timeout=10, environment={'test': 'test'}, get_pty=True, task_id='test_ssh_operator')"
        )

    def test_render_operator_ssh_host_variable(self, ssh_operator_balancer):
        imports, globals_, task = ssh_operator_balancer.render()
        assert imports == [
            "from airflow.providers.ssh.operators.ssh import SSHOperator",
            "from airflow.providers.ssh.hooks.ssh import SSHHook",
            "from airflow.models.variable import Variable",
        ]
        assert globals_ == []
        assert (
            task
            == "SSHOperator(do_xcom_push=True, ssh_hook=SSHHook(remote_host='test_host.local', username='test_user', password=Variable.get('VAR')['password']), ssh_conn_id='test', command='test', cmd_timeout=10, environment={'test': 'test'}, get_pty=True, task_id='test_ssh_operator')"
        )

    def test_render_dag(self, dag):
        assert isinstance(dag, Dag)
        assert (
            dag.render(debug_filename="airflow_pydantic/tests/rendered/dag.py")
            == """# Generated by airflow-config
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.sensors.bash import BashSensor
from airflow.sensors.python import PythonSensor
from airflow_pydantic.tests.conftest import foo
from datetime import datetime
from datetime import timedelta

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    default_view="grid",
    orientation="LR",
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    enabled=True,
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task1 = PythonOperator(
        python_callable=foo,
        op_args=["test"],
        op_kwargs={"test": "test"},
        templates_dict={"test": "test"},
        templates_exts=[".sql", ".hql"],
        show_return_value_in_logs=True,
        task_id="test_python_operator",
        dag=dag,
    )
    task2 = BashOperator(
        bash_command="test",
        env={"test": "test"},
        append_env=True,
        output_encoding="utf-8",
        skip_on_exit_code=99,
        cwd="test",
        output_processor=foo,
        task_id="test_bash_operator",
        dag=dag,
    )
    task3 = SSHOperator(
        do_xcom_push=True,
        ssh_hook=SSHHook(remote_host="test", username="test"),
        ssh_conn_id="test",
        command="test",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="test_ssh_operator",
        dag=dag,
    )
    task4 = BashSensor(bash_command="test", task_id="test_bash_sensor", cwd="test", dag=dag)
    task5 = PythonSensor(python_callable=foo, op_args=["test"], op_kwargs={"test": "test"}, task_id="test_python_sensor", dag=dag)
"""
        )

    def test_render_with_dependencies(self, dag):
        dag.tasks["task1"].dependencies = []
        dag.tasks["task2"].dependencies = ["task1"]
        dag.tasks["task3"].dependencies = ["task1", "task2"]
        assert isinstance(dag, Dag)
        assert (
            dag.render(debug_filename="airflow_pydantic/tests/rendered/dependencies.py")
            == """# Generated by airflow-config
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.providers.ssh.hooks.ssh import SSHHook
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow.sensors.bash import BashSensor
from airflow.sensors.python import PythonSensor
from airflow_pydantic.tests.conftest import foo
from datetime import datetime
from datetime import timedelta

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    default_view="grid",
    orientation="LR",
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    enabled=True,
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task1 = PythonOperator(
        python_callable=foo,
        op_args=["test"],
        op_kwargs={"test": "test"},
        templates_dict={"test": "test"},
        templates_exts=[".sql", ".hql"],
        show_return_value_in_logs=True,
        task_id="test_python_operator",
        dag=dag,
    )
    task2 = BashOperator(
        bash_command="test",
        env={"test": "test"},
        append_env=True,
        output_encoding="utf-8",
        skip_on_exit_code=99,
        cwd="test",
        output_processor=foo,
        task_id="test_bash_operator",
        dag=dag,
    )
    task3 = SSHOperator(
        do_xcom_push=True,
        ssh_hook=SSHHook(remote_host="test", username="test"),
        ssh_conn_id="test",
        command="test",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="test_ssh_operator",
        dag=dag,
    )
    task4 = BashSensor(bash_command="test", task_id="test_bash_sensor", cwd="test", dag=dag)
    task5 = PythonSensor(python_callable=foo, op_args=["test"], op_kwargs={"test": "test"}, task_id="test_python_sensor", dag=dag)
    task1 >> task2
    task1 >> task3
    task2 >> task3
"""
        )

    def test_render_with_externals(self, dag_with_external):
        assert isinstance(dag_with_external, Dag)
        assert (
            dag_with_external.render(debug_filename="airflow_pydantic/tests/rendered/externals.py")
            == """# Generated by airflow-config
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow_pydantic.tests.conftest import foo
from airflow_pydantic.tests.conftest import hook
from datetime import datetime
from datetime import timedelta

with DAG(
    description="",
    schedule="* * * * *",
    start_date=datetime.fromisoformat("2025-01-01T00:00:00"),
    end_date=datetime.fromisoformat("2026-01-01T00:00:00"),
    max_active_tasks=1,
    max_active_runs=1,
    default_view="grid",
    orientation="LR",
    catchup=False,
    is_paused_upon_creation=True,
    tags=["a", "b"],
    dag_display_name="test",
    enabled=True,
    dag_id="a-dag",
    default_args={
        "owner": "airflow",
        "email": ["test@test.com"],
        "email_on_failure": True,
        "email_on_retry": True,
        "retries": 3,
        "retry_delay": timedelta(300.0),
        "start_date": datetime.fromisoformat("2025-01-01T00:00:00"),
        "end_date": datetime.fromisoformat("2026-01-01T00:00:00"),
        "depends_on_past": True,
        "queue": "default",
        "pool": "default",
        "pool_slots": 1,
        "do_xcom_push": True,
        "task_display_name": "test",
    },
) as dag:
    task1 = PythonOperator(
        python_callable=foo,
        op_args=["test"],
        op_kwargs={"test": "test"},
        templates_dict={"test": "test"},
        templates_exts=[".sql", ".hql"],
        show_return_value_in_logs=True,
        task_id="test_python_operator",
        dag=dag,
    )
    task2 = BashOperator(
        bash_command="test",
        env={"test": "test"},
        append_env=True,
        output_encoding="utf-8",
        skip_on_exit_code=99,
        cwd="test",
        output_processor=foo,
        task_id="test_bash_operator",
        dag=dag,
    )
    task3 = SSHOperator(
        do_xcom_push=True,
        ssh_hook=hook(),
        ssh_conn_id="test",
        command="test",
        cmd_timeout=10,
        environment={"test": "test"},
        get_pty=True,
        task_id="test_ssh_operator",
        dag=dag,
    )
"""
        )
