#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import datetime
import importlib.util
import inspect
from unittest import mock
from unittest.mock import MagicMock

import pytest

from airflow import DAG
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.models import Connection
from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.common.sql.hooks.handlers import fetch_all_handler
from airflow.providers.common.sql.hooks.sql import DbApiHook
from airflow.providers.common.sql.operators.sql import (
    BaseSQLOperator,
    BranchSQLOperator,
    SQLCheckOperator,
    SQLCheckResult,
    SQLColumnCheckOperator,
    SQLExecuteQueryOperator,
    SQLInsertRowsOperator,
    SQLIntervalCheckOperator,
    SQLTableCheckOperator,
    SQLThresholdCheckOperator,
    SQLValueCheckOperator,
)
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils import timezone  # type: ignore[attr-defined]
from airflow.utils.state import State
from airflow.utils.types import DagRunType

from tests_common.test_utils.dag import sync_dag_to_db
from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags, clear_db_runs, clear_db_xcom
from tests_common.test_utils.markers import skip_if_force_lowest_dependencies_marker
from tests_common.test_utils.providers import get_provider_min_airflow_version
from tests_common.test_utils.taskinstance import TaskInstanceWrapper
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_1, AIRFLOW_V_3_0_PLUS

try:
    import importlib.util

    if not importlib.util.find_spec("airflow.sdk.bases.hook"):
        raise ImportError

    BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook"
except ImportError:
    BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook"

if AIRFLOW_V_3_0_PLUS:
    from airflow.utils.types import DagRunTriggeredByType

pytestmark = pytest.mark.db_test


class MockHook:
    def get_records(self):
        return


def _get_mock_db_hook():
    return MockHook()


class TestBaseSQLOperator:
    def _construct_operator(self, **kwargs):
        dag = DAG(
            "test_dag",
            schedule=None,
            start_date=datetime.datetime(2017, 1, 1),
            render_template_as_native_obj=True,
        )
        return BaseSQLOperator(
            task_id="test_task",
            conn_id="{{ conn_id }}",
            database="{{ database }}",
            hook_params="{{ hook_params }}",
            **kwargs,
            dag=dag,
        )

    def test_templated_fields(self):
        operator = self._construct_operator()
        operator.render_template_fields(
            {"conn_id": "my_conn_id", "database": "my_database", "hook_params": {"key": "value"}}
        )
        assert operator.conn_id == "my_conn_id"
        assert operator.database == "my_database"
        assert operator.hook_params == {"key": "value"}

    def test_when_provider_min_airflow_version_is_3_0_or_higher_remove_obsolete_get_hook_method(self):
        """
        Once this test starts failing due to the fact that the minimum Airflow version is now 3.0.0 or higher
        for this provider, you should remove the obsolete get_hook method in the BaseSQLOperator operator
        and remove this test.  This test was added to make sure to not forget to remove the fallback code
        for backward compatibility with Airflow 2.8.x which isn't need anymore once this provider depends on
        Airflow 3.0.0 or higher.
        """
        min_airflow_version = get_provider_min_airflow_version("apache-airflow-providers-common-sql")

        # Check if the current Airflow version is 3.0.0 or higher
        if min_airflow_version[0] >= 3:
            method_source = inspect.getsource(BaseSQLOperator.get_hook)
            raise AirflowProviderDeprecationWarning(
                f"Check TODO's to remove obsolete get_hook method in BaseSQLOperator:\n\r\n\r\t\t\t{method_source}"
            )


class TestSQLExecuteQueryOperator:
    def setup_method(self):
        self.task_id = "test_task"
        self.conn_id = "sql_default"
        self._operator = SQLExecuteQueryOperator(task_id=self.task_id, conn_id=self.conn_id, sql="sql")

    def _construct_operator(self, sql, **kwargs):
        dag = DAG("test_dag", schedule=None, start_date=datetime.datetime(2017, 1, 1))
        return SQLExecuteQueryOperator(
            task_id="test_task",
            conn_id="default_conn",
            sql=sql,
            **kwargs,
            dag=dag,
        )

    @mock.patch.object(SQLExecuteQueryOperator, "_process_output")
    @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
    @pytest.mark.parametrize("requires_result_fetch", [True, False])
    def test_do_xcom_push(self, mock_get_db_hook, mock_process_output, requires_result_fetch):
        operator = self._construct_operator(
            "SELECT 1;", do_xcom_push=True, requires_result_fetch=requires_result_fetch
        )
        operator.execute(context=MagicMock())

        mock_get_db_hook.return_value.run.assert_called_once_with(
            sql="SELECT 1;",
            autocommit=False,
            handler=fetch_all_handler,
            parameters=None,
            return_last=True,
        )
        mock_process_output.assert_called()

    @mock.patch.object(SQLExecuteQueryOperator, "_process_output")
    @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
    def test_dont_xcom_push(self, mock_get_db_hook, mock_process_output):
        operator = self._construct_operator("SELECT 1;", do_xcom_push=False)
        operator.execute(context=MagicMock())

        mock_get_db_hook.return_value.run.assert_called_once_with(
            sql="SELECT 1;",
            autocommit=False,
            parameters=None,
            handler=None,
            return_last=True,
        )
        mock_process_output.assert_not_called()

    @mock.patch.object(SQLExecuteQueryOperator, "_process_output")
    @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
    def test_requires_result_fetch_dont_xcom_push(self, mock_get_db_hook, mock_process_output):
        operator = self._construct_operator("SELECT 1;", requires_result_fetch=True, do_xcom_push=False)
        operator.execute(context=MagicMock())

        mock_get_db_hook.return_value.run.assert_called_once_with(
            sql="SELECT 1;",
            autocommit=False,
            handler=fetch_all_handler,
            parameters=None,
            return_last=True,
        )
        mock_process_output.assert_not_called()

    @mock.patch.object(SQLExecuteQueryOperator, "get_db_hook")
    def test_output_processor(self, mock_get_db_hook):
        data = [(1, "Alice"), (2, "Bob")]

        mock_hook = MagicMock()
        mock_hook.run.return_value = data
        mock_hook.descriptions = ("id", "name")
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator(
            sql="SELECT * FROM users;",
            output_processor=lambda results, descriptions: (descriptions, results),
            return_last=False,
        )
        descriptions, result = operator.execute(context=MagicMock())

        assert descriptions == ("id", "name")
        assert result == [(1, "Alice"), (2, "Bob")]

    @skip_if_force_lowest_dependencies_marker
    def test_sql_operator_extra_dejson_fields_to_hook_params(self):
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.BaseHook.get_connection",
            return_value=Connection(conn_id="sql_default", conn_type="postgres"),
        ) as mock_get_conn:
            mock_get_conn.return_value = Connection(
                conn_id="google_cloud_bigquery_default",
                conn_type="gcpbigquery",
                extra={"use_legacy_sql": False, "priority": "INTERACTIVE"},
            )
            self._operator.hook_params = {"use_legacy_sql": True, "location": "us-east1"}
            assert self._operator._hook.conn_type == "gcpbigquery"
            assert self._operator._hook.use_legacy_sql is True
            assert self._operator._hook.location == "us-east1"
            assert self._operator._hook.priority == "INTERACTIVE"


class TestColumnCheckOperator:
    valid_column_mapping = {
        "X": {
            "null_check": {"equal_to": 0},
            "distinct_check": {"equal_to": 10, "tolerance": 0.1},
            "unique_check": {"geq_to": 10},
            "min": {"leq_to": 1},
            "max": {"less_than": 20, "greater_than": 10},
        }
    }

    short_valid_column_mapping = {
        "X": {
            "null_check": {"equal_to": 0},
            "distinct_check": {"equal_to": 10, "tolerance": 0.1},
        }
    }

    invalid_column_mapping = {"Y": {"invalid_check_name": {"expectation": 5}}}

    correct_generate_sql_query_no_partitions = """
        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table ) AS sq
    UNION ALL
        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table ) AS sq
    """

    correct_generate_sql_query_with_partition = """
        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table WHERE Y > 1) AS sq
    UNION ALL
        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Y > 1) AS sq
    """

    correct_generate_sql_query_with_partition_and_where = """
        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table WHERE Y > 1 AND Z < 100) AS sq
    UNION ALL
        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Y > 1) AS sq
    """

    correct_generate_sql_query_with_where = """
        SELECT 'X' AS col_name, 'null_check' AS check_type, X_null_check AS check_result
        FROM (SELECT SUM(CASE WHEN X IS NULL THEN 1 ELSE 0 END) AS X_null_check FROM test_table ) AS sq
    UNION ALL
        SELECT 'X' AS col_name, 'distinct_check' AS check_type, X_distinct_check AS check_result
        FROM (SELECT COUNT(DISTINCT(X)) AS X_distinct_check FROM test_table WHERE Z < 100) AS sq
    """

    def _construct_operator(self, monkeypatch, column_mapping, records):
        def get_records(*arg):
            return records

        operator = SQLColumnCheckOperator(
            task_id="test_task", table="test_table", column_mapping=column_mapping
        )
        monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
        monkeypatch.setattr(MockHook, "get_records", get_records)
        return operator

    def _full_check_sql(self, sql: str) -> str:
        """
        Wraps the check fragment in the outer parts of the sql query
        """
        return f"SELECT col_name, check_type, check_result FROM ({sql}) AS check_columns"

    def test_check_not_in_column_checks(self, monkeypatch):
        with pytest.raises(AirflowException, match="Invalid column check: invalid_check_name."):
            self._construct_operator(monkeypatch, self.invalid_column_mapping, ())

    def test_pass_all_checks_exact_check(self, monkeypatch):
        records = [
            ("X", "null_check", 0),
            ("X", "distinct_check", 10),
            ("X", "unique_check", 10),
            ("X", "min", 1),
            ("X", "max", 19),
        ]
        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
        operator.execute(context=MagicMock())
        assert [
            operator.column_mapping["X"][check]["success"] is True
            for check in [*operator.column_mapping["X"]]
        ]

    def test_max_less_than_fails_check(self, monkeypatch):
        records = [
            ("X", "null_check", 1),
            ("X", "distinct_check", 10),
            ("X", "unique_check", 10),
            ("X", "min", 1),
            ("X", "max", 21),
        ]
        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
        with pytest.raises(AirflowException, match="Test failed") as err_ctx:
            operator.execute(context=MagicMock())
        assert "Check: max" in str(err_ctx.value)
        assert "{'less_than': 20, 'greater_than': 10, 'result': 21, 'success': False}" in str(err_ctx.value)
        assert operator.column_mapping["X"]["max"]["success"] is False

    def test_max_greater_than_fails_check(self, monkeypatch):
        records = [
            ("X", "null_check", 1),
            ("X", "distinct_check", 10),
            ("X", "unique_check", 10),
            ("X", "min", 1),
            ("X", "max", 9),
        ]
        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
        with pytest.raises(AirflowException, match="Test failed") as err_ctx:
            operator.execute(context=MagicMock())
        assert "Check: max" in str(err_ctx.value)
        assert "{'less_than': 20, 'greater_than': 10, 'result': 9, 'success': False}" in str(err_ctx.value)
        assert operator.column_mapping["X"]["max"]["success"] is False

    def test_pass_all_checks_inexact_check(self, monkeypatch):
        records = [
            ("X", "null_check", 0),
            ("X", "distinct_check", 9),
            ("X", "unique_check", 12),
            ("X", "min", 0),
            ("X", "max", 15),
        ]
        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
        operator.execute(context=MagicMock())
        assert [
            operator.column_mapping["X"][check]["success"] is True
            for check in [*operator.column_mapping["X"]]
        ]

    def test_fail_all_checks_check(self, monkeypatch):
        records = [
            ("X", "null_check", 1),
            ("X", "distinct_check", 12),
            ("X", "unique_check", 11),
            ("X", "min", -1),
            ("X", "max", 20),
        ]
        operator = self._construct_operator(monkeypatch, self.valid_column_mapping, records)
        with pytest.raises(AirflowException):
            operator.execute(context=MagicMock())

    def test_generate_sql_query_no_partitions(self, monkeypatch):
        checks = self.short_valid_column_mapping["X"]
        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
        assert (
            operator._generate_sql_query("X", checks).lstrip()
            == self.correct_generate_sql_query_no_partitions.lstrip()
        )

    def test_generate_sql_query_with_partitions(self, monkeypatch):
        checks = self.short_valid_column_mapping["X"]
        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
        operator.partition_clause = "Y > 1"
        assert (
            operator._generate_sql_query("X", checks).lstrip()
            == self.correct_generate_sql_query_with_partition.lstrip()
        )

    def test_generate_sql_query_with_templated_partitions(self, monkeypatch):
        checks = self.short_valid_column_mapping["X"]
        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
        operator.partition_clause = "{{ params.col }} > 1"
        operator.render_template_fields({"params": {"col": "Y"}})
        assert (
            operator._generate_sql_query("X", checks).lstrip()
            == self.correct_generate_sql_query_with_partition.lstrip()
        )

    def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
        self.short_valid_column_mapping["X"]["null_check"]["partition_clause"] = "Z < 100"
        checks = self.short_valid_column_mapping["X"]
        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
        operator.partition_clause = "Y > 1"
        assert (
            operator._generate_sql_query("X", checks).lstrip()
            == self.correct_generate_sql_query_with_partition_and_where.lstrip()
        )
        del self.short_valid_column_mapping["X"]["null_check"]["partition_clause"]

    def test_generate_sql_query_with_check_partition(self, monkeypatch):
        self.short_valid_column_mapping["X"]["distinct_check"]["partition_clause"] = "Z < 100"
        checks = self.short_valid_column_mapping["X"]
        operator = self._construct_operator(monkeypatch, self.short_valid_column_mapping, ())
        assert (
            operator._generate_sql_query("X", checks).lstrip()
            == self.correct_generate_sql_query_with_where.lstrip()
        )
        del self.short_valid_column_mapping["X"]["distinct_check"]["partition_clause"]

    @mock.patch.object(SQLColumnCheckOperator, "get_db_hook")
    def test_generated_sql_respects_templated_partitions(self, mock_get_db_hook):
        records = [
            ("X", "null_check", 0),
            ("X", "distinct_check", 10),
        ]

        mock_hook = mock.Mock()
        mock_hook.get_records.return_value = records
        mock_get_db_hook.return_value = mock_hook

        operator = SQLColumnCheckOperator(
            task_id="test_task",
            table="test_table",
            column_mapping=self.short_valid_column_mapping,
            partition_clause="{{ params.col }} > 1",
        )
        operator.render_template_fields({"params": {"col": "Y"}})

        operator.execute(context=MagicMock())

        mock_get_db_hook.return_value.get_records.assert_called_once_with(
            self._full_check_sql(self.correct_generate_sql_query_with_partition),
        )

    @mock.patch.object(SQLColumnCheckOperator, "get_db_hook")
    def test_generated_sql_respects_templated_table(self, mock_get_db_hook):
        records = [
            ("X", "null_check", 0),
            ("X", "distinct_check", 10),
        ]

        mock_hook = mock.Mock()
        mock_hook.get_records.return_value = records
        mock_get_db_hook.return_value = mock_hook

        operator = SQLColumnCheckOperator(
            task_id="test_task",
            table="{{ params.table }}",
            column_mapping=self.short_valid_column_mapping,
        )
        operator.render_template_fields({"params": {"table": "test_table"}})

        operator.execute(context=MagicMock())

        mock_get_db_hook.return_value.get_records.assert_called_once_with(
            self._full_check_sql(self.correct_generate_sql_query_no_partitions),
        )


class TestTableCheckOperator:
    count_check = "COUNT(*) == 1000"
    sum_check = "col_a + col_b < col_c"
    checks = {
        "row_count_check": {"check_statement": f"{count_check}"},
        "column_sum_check": {"check_statement": f"{sum_check}"},
    }

    correct_generate_sql_query_no_partitions = f"""
    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
          FROM test_table ) AS sq
    UNION ALL
    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
          FROM test_table ) AS sq
    """

    correct_generate_sql_query_with_partition = f"""
    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
          FROM test_table WHERE col_a > 10) AS sq
    UNION ALL
    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
          FROM test_table WHERE col_a > 10) AS sq
    """

    correct_generate_sql_query_with_partition_and_where = f"""
    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
          FROM test_table WHERE col_a > 10 AND id = 100) AS sq
    UNION ALL
    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
          FROM test_table WHERE col_a > 10) AS sq
    """

    correct_generate_sql_query_with_where = f"""
    SELECT 'row_count_check' AS check_name, MIN(row_count_check) AS check_result
    FROM (SELECT CASE WHEN {count_check} THEN 1 ELSE 0 END AS row_count_check
          FROM test_table ) AS sq
    UNION ALL
    SELECT 'column_sum_check' AS check_name, MIN(column_sum_check) AS check_result
    FROM (SELECT CASE WHEN {sum_check} THEN 1 ELSE 0 END AS column_sum_check
          FROM test_table WHERE id = 100) AS sq
    """

    def _construct_operator(self, monkeypatch, checks, records):
        def get_records(*arg):
            return records

        operator = SQLTableCheckOperator(task_id="test_task", table="test_table", checks=checks)
        monkeypatch.setattr(operator, "get_db_hook", _get_mock_db_hook)
        monkeypatch.setattr(MockHook, "get_records", get_records)
        return operator

    @pytest.mark.parametrize(
        "conn_id",
        [
            pytest.param("postgres_default", marks=[pytest.mark.backend("postgres")]),
            pytest.param("mysql_default", marks=[pytest.mark.backend("mysql")]),
        ],
    )
    def test_sql_check(self, conn_id):
        operator = SQLTableCheckOperator(
            task_id="test_task",
            table="employees",
            checks={"row_count_check": {"check_statement": "COUNT(*) >= 3"}},
            conn_id=conn_id,
        )

        hook = operator.get_db_hook()
        hook.run(
            [
                """
                CREATE TABLE IF NOT EXISTS employees (
                    employee_name VARCHAR(50) NOT NULL,
                    employment_year INT NOT NULL
                );
                """,
                "INSERT INTO employees VALUES ('Adam', 2021)",
                "INSERT INTO employees VALUES ('Chris', 2021)",
                "INSERT INTO employees VALUES ('Frank', 2021)",
                "INSERT INTO employees VALUES ('Fritz', 2021)",
                "INSERT INTO employees VALUES ('Magda', 2022)",
                "INSERT INTO employees VALUES ('Phil', 2021)",
            ]
        )
        try:
            operator.execute({})
        finally:
            hook.run(["DROP TABLE employees"])

    @pytest.mark.parametrize(
        "conn_id",
        [
            pytest.param("postgres_default", marks=[pytest.mark.backend("postgres")]),
            pytest.param("mysql_default", marks=[pytest.mark.backend("mysql")]),
        ],
    )
    def test_sql_check_partition_clause_templating(self, conn_id):
        """
        Checks that the generated sql respects a templated partition clause
        """
        operator = SQLTableCheckOperator(
            task_id="test_task",
            table="employees",
            checks={"row_count_check": {"check_statement": "COUNT(*) = 5"}},
            conn_id=conn_id,
            partition_clause="employment_year = {{ params.year }}",
        )

        hook = operator.get_db_hook()
        hook.run(
            [
                """
                CREATE TABLE IF NOT EXISTS employees (
                    employee_name VARCHAR(50) NOT NULL,
                    employment_year INT NOT NULL
                );
                """,
                "INSERT INTO employees VALUES ('Adam', 2021)",
                "INSERT INTO employees VALUES ('Chris', 2021)",
                "INSERT INTO employees VALUES ('Frank', 2021)",
                "INSERT INTO employees VALUES ('Fritz', 2021)",
                "INSERT INTO employees VALUES ('Magda', 2022)",
                "INSERT INTO employees VALUES ('Phil', 2021)",
            ]
        )
        try:
            operator.render_template_fields({"params": {"year": 2021}})
            operator.execute({})
        finally:
            hook.run(["DROP TABLE employees"])

    def test_sql_check_accept_none_true(self, monkeypatch):
        """Validate that empty table does not fail when accept_none=True."""
        records = []
        operator = self._construct_operator(monkeypatch, self.checks, records)
        operator.accept_none = True
        operator.execute(context=MagicMock())

    def test_sql_check_accept_none_false(self, monkeypatch):
        """Validate that empty table throws an exception when accept_none=False."""
        records = []
        operator = self._construct_operator(monkeypatch, self.checks, records)
        operator.accept_none = False  # This is default, technically
        with pytest.raises(AirflowException):
            operator.execute(context=MagicMock())

    def test_pass_all_checks_check(self, monkeypatch):
        records = [("row_count_check", 1), ("column_sum_check", "y")]
        operator = self._construct_operator(monkeypatch, self.checks, records)
        operator.execute(context=MagicMock())
        assert [operator.checks[check]["success"] is True for check in operator.checks.keys()]

    def test_fail_all_checks_check(self, monkeypatch):
        records = [("row_count_check", 0), ("column_sum_check", "n")]
        operator = self._construct_operator(monkeypatch, self.checks, records)
        with pytest.raises(AirflowException):
            operator.execute(context=MagicMock())

    def test_generate_sql_query_no_partitions(self, monkeypatch):
        operator = self._construct_operator(monkeypatch, self.checks, ())
        assert (
            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_no_partitions.lstrip()
        )

    def test_generate_sql_query_with_partitions(self, monkeypatch):
        operator = self._construct_operator(monkeypatch, self.checks, ())
        operator.partition_clause = "col_a > 10"
        assert (
            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_partition.lstrip()
        )

    def test_generate_sql_query_with_templated_partitions(self, monkeypatch):
        operator = self._construct_operator(monkeypatch, self.checks, ())
        operator.partition_clause = "{{ params.col }} > 10"
        operator.render_template_fields({"params": {"col": "col_a"}})
        assert (
            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_partition.lstrip()
        )

    def test_generate_sql_query_with_templated_table(self, monkeypatch):
        operator = self._construct_operator(monkeypatch, self.checks, ())
        operator.table = "{{ params.table }}"
        operator.render_template_fields({"params": {"table": "test_table"}})
        assert (
            operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_no_partitions.lstrip()
        )

    def test_generate_sql_query_with_partitions_and_check_partition(self, monkeypatch):
        self.checks["row_count_check"]["partition_clause"] = "id = 100"
        operator = self._construct_operator(monkeypatch, self.checks, ())
        operator.partition_clause = "col_a > 10"
        assert (
            operator._generate_sql_query().lstrip()
            == self.correct_generate_sql_query_with_partition_and_where.lstrip()
        )
        del self.checks["row_count_check"]["partition_clause"]

    def test_generate_sql_query_with_check_partition(self, monkeypatch):
        self.checks["column_sum_check"]["partition_clause"] = "id = 100"
        operator = self._construct_operator(monkeypatch, self.checks, ())
        assert operator._generate_sql_query().lstrip() == self.correct_generate_sql_query_with_where.lstrip()
        del self.checks["column_sum_check"]["partition_clause"]


DEFAULT_DATE = timezone.datetime(2016, 1, 1)
INTERVAL = datetime.timedelta(hours=12)
SUPPORTED_TRUE_VALUES = [
    ["True"],
    ["true"],
    ["1"],
    ["on"],
    [1],
    True,
    "true",
    "1",
    "on",
    1,
]
SUPPORTED_FALSE_VALUES = [
    ["False"],
    ["false"],
    ["0"],
    ["off"],
    [0],
    False,
    "false",
    "0",
    "off",
    0,
]


class TestSQLCheckOperatorDbHook:
    def setup_method(self):
        self.task_id = "test_task"
        self.conn_id = "sql_default"
        self._operator = SQLCheckOperator(task_id=self.task_id, conn_id=self.conn_id, sql="sql")

    @pytest.mark.parametrize("database", [None, "test-db"])
    def test_get_hook(self, database):
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.BaseHook.get_connection",
            return_value=Connection(conn_id="sql_default", conn_type="postgres"),
        ) as mock_get_conn:
            if database:
                self._operator.database = database
            assert isinstance(self._operator._hook, PostgresHook)
            mock_get_conn.assert_called_once_with(self.conn_id)

    @skip_if_force_lowest_dependencies_marker
    def test_not_allowed_conn_type(self):
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.BaseHook.get_connection",
            return_value=Connection(conn_id="sql_default", conn_type="postgres"),
        ) as mock_get_conn:
            mock_get_conn.return_value = Connection(conn_id="sql_default", conn_type="airbyte")
            with pytest.raises(AirflowException, match=r"You are trying to use `common-sql`"):
                self._operator._hook

    @skip_if_force_lowest_dependencies_marker
    def test_sql_operator_hook_params_snowflake(self):
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.BaseHook.get_connection",
            return_value=Connection(conn_id="sql_default", conn_type="postgres"),
        ) as mock_get_conn:
            mock_get_conn.return_value = Connection(conn_id="snowflake_default", conn_type="snowflake")
            self._operator.hook_params = {
                "warehouse": "warehouse",
                "database": "database",
                "role": "role",
                "schema": "schema",
                "log_sql": False,
            }
            assert self._operator._hook.conn_type == "snowflake"
            assert self._operator._hook.warehouse == "warehouse"
            assert self._operator._hook.database == "database"
            assert self._operator._hook.role == "role"
            assert self._operator._hook.schema == "schema"
            assert not self._operator._hook.log_sql

    @skip_if_force_lowest_dependencies_marker
    def test_sql_operator_hook_params_biguery(self):
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.BaseHook.get_connection",
            return_value=Connection(conn_id="sql_default", conn_type="postgres"),
        ) as mock_get_conn:
            mock_get_conn.return_value = Connection(
                conn_id="google_cloud_bigquery_default", conn_type="gcpbigquery"
            )
            self._operator.hook_params = {"use_legacy_sql": True, "location": "us-east1"}
            assert self._operator._hook.conn_type == "gcpbigquery"
            assert self._operator._hook.use_legacy_sql
            assert self._operator._hook.location == "us-east1"

    @skip_if_force_lowest_dependencies_marker
    def test_sql_operator_hook_params_templated(self):
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.BaseHook.get_connection",
            return_value=Connection(conn_id="sql_default", conn_type="postgres"),
        ) as mock_get_conn:
            mock_get_conn.return_value = Connection(conn_id="snowflake_default", conn_type="snowflake")
            self._operator.hook_params = {"session_parameters": {"query_tag": "{{ ds }}"}}
            logical_date = "2024-04-02"
            self._operator.render_template_fields({"ds": logical_date})

            assert self._operator._hook.conn_type == "snowflake"
            assert self._operator._hook.session_parameters == {"query_tag": logical_date}


class TestCheckOperator:
    def setup_method(self):
        self._operator = SQLCheckOperator(task_id="test_task", sql="sql", parameters="parameters")

    @mock.patch.object(SQLCheckOperator, "get_db_hook")
    def test_execute_no_records(self, mock_get_db_hook):
        mock_get_db_hook.return_value.get_first.return_value = []

        with pytest.raises(AirflowException, match=r"The following query returned zero rows: sql"):
            self._operator.execute({})

    @mock.patch.object(SQLCheckOperator, "get_db_hook")
    def test_execute_not_all_records_are_true(self, mock_get_db_hook):
        mock_get_db_hook.return_value.get_first.return_value = ["data", ""]

        with pytest.raises(AirflowException, match=r"Test failed."):
            self._operator.execute({})

    @mock.patch.object(SQLCheckOperator, "get_db_hook")
    def test_execute_records_dict_not_all_values_are_true(self, mock_get_db_hook):
        mock_get_db_hook.return_value.get_first.return_value = {
            "DUPLICATE_ID_CHECK": False,
            "NULL_VALUES_CHECK": True,
        }

        with pytest.raises(AirflowException, match=r"Test failed."):
            self._operator.execute({})

    @mock.patch.object(SQLCheckOperator, "get_db_hook")
    def test_sqlcheckoperator_parameters(self, mock_get_db_hook):
        self._operator.execute({})
        mock_get_db_hook.return_value.get_first.assert_called_once_with("sql", "parameters")


class TestValueCheckOperator:
    def setup_method(self):
        self.task_id = "test_task"
        self.conn_id = "default_conn"

    def _construct_operator(self, sql, pass_value, tolerance=None):
        dag = DAG("test_dag", schedule=None, start_date=datetime.datetime(2017, 1, 1))

        return SQLValueCheckOperator(
            dag=dag,
            task_id=self.task_id,
            conn_id=self.conn_id,
            sql=sql,
            pass_value=pass_value,
            tolerance=tolerance,
        )

    def test_pass_value_template_string(self):
        pass_value_str = "2018-03-22"
        operator = self._construct_operator("select date from tab1;", "{{ ds }}")

        operator.render_template_fields({"ds": pass_value_str})

        assert operator.task_id == self.task_id
        assert operator.pass_value == pass_value_str

    def test_pass_value_template_string_float(self):
        pass_value_float = 4.0
        operator = self._construct_operator("select date from tab1;", pass_value_float)

        operator.render_template_fields({})

        assert operator.task_id == self.task_id
        assert operator.pass_value == str(pass_value_float)

    @mock.patch.object(SQLValueCheckOperator, "get_db_hook")
    def test_execute_pass(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = [10]
        mock_get_db_hook.return_value = mock_hook
        sql = "select value from tab1 limit 1;"
        operator = self._construct_operator(sql, 5, 1)

        operator.execute(None)

        mock_hook.get_first.assert_called_once_with(sql, None)

    @mock.patch.object(SQLValueCheckOperator, "get_db_hook")
    def test_execute_fail(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = [11]
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("select value from tab1 limit 1;", 5, 1)

        with pytest.raises(AirflowException, match="Tolerance:100.0%"):
            operator.execute(context=MagicMock())


class TestIntervalCheckOperator:
    def _construct_operator(self, table, metric_thresholds, ratio_formula, ignore_zero):
        return SQLIntervalCheckOperator(
            task_id="test_task",
            table=table,
            metrics_thresholds=metric_thresholds,
            ratio_formula=ratio_formula,
            ignore_zero=ignore_zero,
        )

    def test_invalid_ratio_formula(self):
        with pytest.raises(AirflowException, match="Invalid diff_method"):
            self._construct_operator(
                table="test_table",
                metric_thresholds={
                    "f1": 1,
                },
                ratio_formula="abs",
                ignore_zero=False,
            )

    @mock.patch.object(SQLIntervalCheckOperator, "get_db_hook")
    def test_execute_not_ignore_zero(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = [0]
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator(
            table="test_table",
            metric_thresholds={
                "f1": 1,
            },
            ratio_formula="max_over_min",
            ignore_zero=False,
        )

        with pytest.raises(AirflowException):
            operator.execute(context=MagicMock())

    @mock.patch.object(SQLIntervalCheckOperator, "get_db_hook")
    def test_execute_ignore_zero(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = [0]
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator(
            table="test_table",
            metric_thresholds={
                "f1": 1,
            },
            ratio_formula="max_over_min",
            ignore_zero=True,
        )

        operator.execute(context=MagicMock())

    @mock.patch.object(SQLIntervalCheckOperator, "get_db_hook")
    def test_execute_min_max(self, mock_get_db_hook):
        mock_hook = mock.Mock()

        def returned_row():
            rows = [
                [2, 2, 2, 2],  # reference
                [1, 1, 1, 1],  # current
            ]

            yield from rows

        mock_hook.get_first.side_effect = returned_row()
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator(
            table="test_table",
            metric_thresholds={
                "f0": 1.0,
                "f1": 1.5,
                "f2": 2.0,
                "f3": 2.5,
            },
            ratio_formula="max_over_min",
            ignore_zero=True,
        )

        expected_err_message = (
            "The following tests have failed:\n "
            "f0: {'metric': 'f0', 'current_metric': 1, 'past_metric': 2, 'threshold': 1.0,"
            " 'ignore_zero': True, 'ratio': 2.0, 'success': False}; "
            "f1: {'metric': 'f1', 'current_metric': 1, 'past_metric': 2, 'threshold': 1.5,"
            " 'ignore_zero': True, 'ratio': 2.0, 'success': False}; "
            "f2: {'metric': 'f2', 'current_metric': 1, 'past_metric': 2, 'threshold': 2.0,"
            " 'ignore_zero': True, 'ratio': 2.0, 'success': False}"
        )
        with pytest.raises(AirflowException, match=expected_err_message):
            operator.execute(context=MagicMock())

    @mock.patch.object(SQLIntervalCheckOperator, "get_db_hook")
    def test_execute_diff(self, mock_get_db_hook):
        mock_hook = mock.Mock()

        def returned_row():
            rows = [
                [3, 3, 3, 3],  # reference
                [1, 1, 1, 1],  # current
            ]

            yield from rows

        mock_hook.get_first.side_effect = returned_row()
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator(
            table="test_table",
            metric_thresholds={
                "f0": 0.5,
                "f1": 0.6,
                "f2": 0.7,
                "f3": 0.8,
            },
            ratio_formula="relative_diff",
            ignore_zero=True,
        )

        expected_err_message = (
            "The following tests have failed:\n "
            "f0: {'metric': 'f0', 'current_metric': 1, 'past_metric': 3, 'threshold': 0.5, "
            "'ignore_zero': True, 'ratio': 0.6666666666666666, 'success': False}; "
            "f1: {'metric': 'f1', 'current_metric': 1, 'past_metric': 3, 'threshold': 0.6, "
            "'ignore_zero': True, 'ratio': 0.6666666666666666, 'success': False}"
        )
        with pytest.raises(AirflowException, match=expected_err_message):
            operator.execute(context=MagicMock())


class TestThresholdCheckOperator:
    def _construct_operator(self, sql, min_threshold, max_threshold):
        dag = DAG("test_dag", schedule=None, start_date=datetime.datetime(2017, 1, 1))

        return SQLThresholdCheckOperator(
            task_id="test_task",
            sql=sql,
            min_threshold=min_threshold,
            max_threshold=max_threshold,
            dag=dag,
        )

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_pass_min_value_max_value(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = (10,)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("Select avg(val) from table1 limit 1", 1, 100)

        operator.execute(context=MagicMock())

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_pass_min_value_max_value_templated(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = (10,)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("Select avg(val) from table1 limit 1", "{{ params.min }}", 100)
        operator.render_template_fields({"params": {"min": 1}})
        operator.execute(context=MagicMock())
        mock_hook.get_first.assert_called_once_with("Select avg(val) from table1 limit 1")

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_fail_min_value_max_value(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = (10,)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("Select avg(val) from table1 limit 1", 20, 100)

        with pytest.raises(AirflowException, match="10.*20.0.*100.0"):
            operator.execute(context=MagicMock())

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_pass_min_sql_max_sql(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.side_effect = lambda x: (int(x.split()[1]),)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("Select 10", "Select 1", "Select 100")

        operator.execute(context=MagicMock())

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_fail_min_sql_max_sql(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.side_effect = lambda x: (int(x.split()[1]),)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("Select 10", "Select 20", "Select 100")

        with pytest.raises(AirflowException, match="10.*20.*100"):
            operator.execute(context=MagicMock())

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    @pytest.mark.parametrize(
        ("sql", "min_threshold", "max_threshold"),
        (
            ("Select 75", 45, "Select 100"),
            # check corner-case if result of query is "falsey" does not raise error
            ("Select 0", 0, 1),
            ("Select 1", 0, 1),
        ),
    )
    def test_pass_min_value_max_sql(self, mock_get_db_hook, sql, min_threshold, max_threshold):
        mock_hook = mock.Mock()
        mock_hook.get_first.side_effect = lambda x: (int(x.split()[1]),)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator(sql, min_threshold, max_threshold)

        operator.execute(context=MagicMock())

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_fail_min_sql_max_value(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.side_effect = lambda x: (int(x.split()[1]),)
        mock_get_db_hook.return_value = mock_hook

        operator = self._construct_operator("Select 155", "Select 45", 100)

        with pytest.raises(AirflowException, match="155.*45.*100.0"):
            operator.execute(context=MagicMock())

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_fail_if_query_returns_no_rows(self, mock_get_db_hook):
        mock_hook = mock.Mock()
        mock_hook.get_first.return_value = None
        mock_get_db_hook.return_value = mock_hook

        sql = "Select val from table1 where val = 'val not in table'"
        operator = self._construct_operator(sql, 20, 100)

        with pytest.raises(AirflowException, match=f"The following query returned zero rows: {sql}"):
            operator.execute(context=MagicMock())


class TestSqlBranch:
    """
    Test for SQL Branch Operator
    """

    @staticmethod
    def clear_db():
        clear_db_dags()
        clear_db_runs()
        clear_db_xcom()
        if AIRFLOW_V_3_0_PLUS:
            clear_db_dag_bundles()

    @classmethod
    def setup_class(cls):
        cls.clear_db()

    def setup_method(self):
        self.clear_db()

        self.dag = DAG(
            "sql_branch_operator_test",
            default_args={"owner": "airflow", "start_date": DEFAULT_DATE},
            schedule=INTERVAL,
        )
        self.branch_1 = EmptyOperator(task_id="branch_1", dag=self.dag)
        self.branch_2 = EmptyOperator(task_id="branch_2", dag=self.dag)
        self.branch_3 = None
        if AIRFLOW_V_3_0_PLUS:
            self.scheduler_dag = sync_dag_to_db(self.dag)
        else:
            self.scheduler_dag = self.dag

    def get_ti(self, task_id, dr=None):
        if dr is None:
            if AIRFLOW_V_3_0_PLUS:
                dagrun_kwargs = {
                    "logical_date": DEFAULT_DATE,
                    "run_after": DEFAULT_DATE,
                    "triggered_by": DagRunTriggeredByType.TEST,
                }
            else:
                dagrun_kwargs = {"execution_date": DEFAULT_DATE}
            dr = self.scheduler_dag.create_dagrun(
                run_id=f"manual__{timezone.utcnow().isoformat()}",
                run_type=DagRunType.MANUAL,
                start_date=timezone.utcnow(),
                state=State.RUNNING,
                data_interval=(DEFAULT_DATE, DEFAULT_DATE),
                **dagrun_kwargs,
            )

        ti = dr.get_task_instance(task_id)
        task = self.dag.get_task(ti.task_id)

        if AIRFLOW_V_3_0_PLUS:
            ti.task = self.scheduler_dag.get_task(ti.task_id)
            ti = TaskInstanceWrapper(ti, task)
        else:
            ti.task = task

        return ti

    def teardown_method(self):
        self.clear_db()

    @pytest.fixture
    def branch_op(self):
        branch_op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="mysql_default",
            sql="SELECT 1",
            follow_task_ids_if_true=["branch_1"],
            follow_task_ids_if_false=["branch_2"],
            dag=self.dag,
        )
        if AIRFLOW_V_3_0_PLUS:
            self.scheduler_dag = sync_dag_to_db(self.dag)
        return branch_op

    def test_unsupported_conn_type(self):
        """Check if BranchSQLOperator throws an exception for unsupported connection type"""
        op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="redis_default",
            sql="SELECT count(1) FROM INFORMATION_SCHEMA.TABLES",
            follow_task_ids_if_true=["branch_1"],
            follow_task_ids_if_false=["branch_2"],
            dag=self.dag,
        )
        with pytest.raises(AirflowException):
            op.execute({})

    def test_invalid_conn(self, sdk_connection_not_found):
        """Check if BranchSQLOperator throws an exception for invalid connection"""
        op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="invalid_connection",
            sql="SELECT count(1) FROM INFORMATION_SCHEMA.TABLES",
            follow_task_ids_if_true=["branch_1"],
            follow_task_ids_if_false=["branch_2"],
            dag=self.dag,
        )
        with pytest.raises(AirflowException):
            op.execute({})

    def test_invalid_follow_task_true(self, sdk_connection_not_found):
        """Check if BranchSQLOperator throws an exception for invalid connection"""
        op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="invalid_connection",
            sql="SELECT count(1) FROM INFORMATION_SCHEMA.TABLES",
            follow_task_ids_if_true=[],
            follow_task_ids_if_false=["branch_2"],
            dag=self.dag,
        )
        with pytest.raises(AirflowException):
            op.execute({})

    def test_invalid_follow_task_false(self, sdk_connection_not_found):
        """Check if BranchSQLOperator throws an exception for invalid connection"""
        op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="invalid_connection",
            sql="SELECT count(1) FROM INFORMATION_SCHEMA.TABLES",
            follow_task_ids_if_true=["branch_1"],
            follow_task_ids_if_false=[],
            dag=self.dag,
        )
        with pytest.raises(AirflowException):
            op.execute({})

    @pytest.mark.backend("mysql")
    def test_sql_branch_operator_mysql(self, branch_op):
        """Check if BranchSQLOperator works with backend"""

        branch_op.execute({"ti": mock.MagicMock(task=branch_op)})

    @pytest.mark.backend("postgres")
    def test_sql_branch_operator_postgres(self):
        """Check if BranchSQLOperator works with backend"""
        branch_op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="postgres_default",
            sql="SELECT 1",
            follow_task_ids_if_true=["branch_1"],
            follow_task_ids_if_false=["branch_2"],
            dag=self.dag,
        )
        if AIRFLOW_V_3_0_PLUS:
            self.scheduler_dag = sync_dag_to_db(self.dag)
        self.get_ti(branch_op.task_id).run()

    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_branch_single_value_with_dag_run(self, mock_get_db_hook, branch_op):
        """Check BranchSQLOperator branch operation"""
        self.branch_1.set_upstream(branch_op)
        self.branch_2.set_upstream(branch_op)
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        dr = self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first
        mock_get_records.return_value = 1

        if AIRFLOW_V_3_0_1:
            from airflow.providers.common.compat.sdk import DownstreamTasksSkipped

            with pytest.raises(DownstreamTasksSkipped) as exc_info:
                branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)

            assert exc_info.value.tasks == [("branch_2", -1)]
        else:
            self.get_ti(branch_op.task_id, dr).run()
            tis = dr.get_task_instances()

            for ti in tis:
                if ti.task_id == "make_choice":
                    assert ti.state == State.SUCCESS
                elif ti.task_id == "branch_1":
                    assert ti.state == State.NONE
                elif ti.task_id == "branch_2":
                    assert ti.state == State.SKIPPED
                else:
                    raise ValueError(f"Invalid task id {ti.task_id} found!")

    @pytest.mark.parametrize("true_value", SUPPORTED_TRUE_VALUES)
    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_branch_true_with_dag_run(self, mock_get_db_hook, true_value, branch_op):
        """Check BranchSQLOperator branch operation"""
        self.branch_1.set_upstream(branch_op)
        self.branch_2.set_upstream(branch_op)
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        dr = self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first
        mock_get_records.return_value = true_value

        if AIRFLOW_V_3_0_1:
            from airflow.providers.common.compat.sdk import DownstreamTasksSkipped

            with pytest.raises(DownstreamTasksSkipped) as exc_info:
                branch_op.execute({})

            assert exc_info.value.tasks == [("branch_2", -1)]
        else:
            self.get_ti(branch_op.task_id, dr).run()
            tis = dr.get_task_instances()
            for ti in tis:
                if ti.task_id == "make_choice":
                    assert ti.state == State.SUCCESS
                elif ti.task_id == "branch_1":
                    assert ti.state == State.NONE
                elif ti.task_id == "branch_2":
                    assert ti.state == State.SKIPPED
                else:
                    raise ValueError(f"Invalid task id {ti.task_id} found!")

    @pytest.mark.parametrize("false_value", SUPPORTED_FALSE_VALUES)
    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_branch_false_with_dag_run(self, mock_get_db_hook, false_value, branch_op):
        """Check BranchSQLOperator branch operation"""
        self.branch_1.set_upstream(branch_op)
        self.branch_2.set_upstream(branch_op)
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        dr = self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first
        mock_get_records.return_value = false_value

        if AIRFLOW_V_3_0_1:
            from airflow.providers.common.compat.sdk import DownstreamTasksSkipped

            with pytest.raises(DownstreamTasksSkipped) as exc_info:
                branch_op.execute({})
            assert exc_info.value.tasks == [("branch_1", -1)]
        else:
            self.get_ti(branch_op.task_id, dr).run()
            tis = dr.get_task_instances()

            for ti in tis:
                if ti.task_id == "make_choice":
                    assert ti.state == State.SUCCESS
                elif ti.task_id == "branch_1":
                    assert ti.state == State.SKIPPED
                elif ti.task_id == "branch_2":
                    assert ti.state == State.NONE
                else:
                    raise ValueError(f"Invalid task id {ti.task_id} found!")

    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_branch_list_with_dag_run(self, mock_get_db_hook):
        """Checks if the BranchSQLOperator supports branching off to a list of tasks."""
        branch_op = BranchSQLOperator(
            task_id="make_choice",
            conn_id="mysql_default",
            sql="SELECT 1",
            follow_task_ids_if_true=["branch_1", "branch_2"],
            follow_task_ids_if_false=["branch_3"],
            dag=self.dag,
        )

        self.branch_1.set_upstream(branch_op)
        self.branch_2.set_upstream(branch_op)
        self.branch_3 = EmptyOperator(task_id="branch_3", dag=self.dag)
        self.branch_3.set_upstream(branch_op)
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            self.scheduler_dag = sync_dag_to_db(self.dag)
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        dr = self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first
        mock_get_records.return_value = [["1"]]

        if AIRFLOW_V_3_0_1:
            from airflow.providers.common.compat.sdk import DownstreamTasksSkipped

            with pytest.raises(DownstreamTasksSkipped) as exc_info:
                branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
            assert exc_info.value.tasks == [("branch_3", -1)]
        else:
            self.get_ti(branch_op.task_id, dr).run()
            tis = dr.get_task_instances()
            for ti in tis:
                if ti.task_id == "make_choice":
                    assert ti.state == State.SUCCESS
                elif ti.task_id in ("branch_1", "branch_2"):
                    assert ti.state == State.NONE
                elif ti.task_id == "branch_3":
                    assert ti.state == State.SKIPPED
                else:
                    raise ValueError(f"Invalid task id {ti.task_id} found!")

    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_invalid_query_result_with_dag_run(self, mock_get_db_hook, branch_op):
        """Check BranchSQLOperator branch operation"""
        self.branch_1.set_upstream(branch_op)
        self.branch_2.set_upstream(branch_op)
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first

        mock_get_records.return_value = ["Invalid Value"]

        with pytest.raises(AirflowException):
            branch_op.execute({})

    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_with_skip_in_branch_downstream_dependencies(self, mock_get_db_hook, branch_op):
        """Test SQL Branch with skipping all downstream dependencies"""
        branch_op >> self.branch_1 >> self.branch_2
        branch_op >> self.branch_2
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        dr = self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first

        for true_value in SUPPORTED_TRUE_VALUES:
            mock_get_records.return_value = [true_value]

            self.get_ti(branch_op.task_id, dr).run()

            tis = dr.get_task_instances()
            for ti in tis:
                if ti.task_id == "make_choice":
                    assert ti.state == State.SUCCESS
                elif ti.task_id in ("branch_1", "branch_2"):
                    assert ti.state == State.NONE
                else:
                    raise ValueError(f"Invalid task id {ti.task_id} found!")

    @pytest.mark.parametrize("false_value", SUPPORTED_FALSE_VALUES)
    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_with_skip_in_branch_downstream_dependencies2(self, mock_get_db_hook, false_value, branch_op):
        """Test skipping downstream dependency for false condition"""
        branch_op >> self.branch_1 >> self.branch_2
        branch_op >> self.branch_2
        self.scheduler_dag.clear()

        if AIRFLOW_V_3_0_PLUS:
            dagrun_kwargs = {
                "logical_date": DEFAULT_DATE,
                "run_after": DEFAULT_DATE,
                "triggered_by": DagRunTriggeredByType.TEST,
            }
        else:
            dagrun_kwargs = {"execution_date": DEFAULT_DATE}
        dr = self.scheduler_dag.create_dagrun(
            run_id="manual__",
            run_type=DagRunType.MANUAL,
            start_date=timezone.utcnow(),
            state=State.RUNNING,
            data_interval=(DEFAULT_DATE, DEFAULT_DATE),
            **dagrun_kwargs,
        )

        mock_get_records = mock_get_db_hook.return_value.get_first
        mock_get_records.return_value = [false_value]

        if AIRFLOW_V_3_0_1:
            from airflow.providers.common.compat.sdk import DownstreamTasksSkipped

            with pytest.raises(DownstreamTasksSkipped) as exc_info:
                branch_op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
            assert exc_info.value.tasks == [("branch_1", -1)]
        else:
            self.get_ti(branch_op.task_id, dr).run()
            tis = dr.get_task_instances()
            for ti in tis:
                if ti.task_id == "make_choice":
                    assert ti.state == State.SUCCESS
                elif ti.task_id == "branch_1":
                    assert ti.state == State.SKIPPED
                elif ti.task_id == "branch_2":
                    assert ti.state == State.NONE
                else:
                    raise ValueError(f"Invalid task id {ti.task_id} found!")


class TestBaseSQLOperatorSubClass:
    from airflow.providers.common.sql.operators.sql import BaseSQLOperator

    class NewStyleBaseSQLOperatorSubClass(BaseSQLOperator):
        """New style subclass of BaseSQLOperator"""

        conn_id_field = "custom_conn_id_field"

        def __init__(self, custom_conn_id_field="test_conn", **kwargs):
            super().__init__(**kwargs)
            self.custom_conn_id_field = custom_conn_id_field

    class OldStyleBaseSQLOperatorSubClass(BaseSQLOperator):
        """Old style subclass of BaseSQLOperator"""

        def __init__(self, custom_conn_id_field="test_conn", **kwargs):
            super().__init__(conn_id=custom_conn_id_field, **kwargs)

    @pytest.mark.parametrize(
        "operator_class", [NewStyleBaseSQLOperatorSubClass, OldStyleBaseSQLOperatorSubClass]
    )
    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
    def test_new_style_subclass(self, mock_get_connection, operator_class):
        from airflow.providers.common.sql.hooks.sql import DbApiHook

        op = operator_class(task_id="test_task")
        mock_get_connection.return_value.get_hook.return_value = MagicMock(spec=DbApiHook)
        op.get_db_hook()
        mock_get_connection.assert_called_once_with("test_conn")


class TestSQLColumnCheckOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(column_mapping, **kwargs):
        return SQLColumnCheckOperator(
            task_id="test_task", table="test_table", column_mapping=column_mapping, **kwargs
        )

    def test_not_null_check(self):
        op = self._make_operator({"col": {"null_check": {"equal_to": 0}}})
        op.column_mapping["col"]["null_check"]["result"] = 0
        op.column_mapping["col"]["null_check"]["success"] = True
        results = op._build_check_results()
        assert len(results) == 1
        r = results[0]
        assert r.name == "col.null_check"
        assert r.check_type == "not_null"
        assert r.success is True
        assert r.column == "col"
        assert r.table == "test_table"
        assert r.expected == "0"
        assert r.actual == "0"
        assert r.content == "SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)"
        assert r.description == "Column-level statistical check against the configured threshold"
        assert r.severity == "error"
        assert r.params == {"equal_to": 0, "accept_none": True}

    def test_unique_check(self):
        op = self._make_operator({"col": {"unique_check": {"equal_to": 0}}})
        op.column_mapping["col"]["unique_check"]["result"] = 0
        op.column_mapping["col"]["unique_check"]["success"] = True
        results = op._build_check_results()
        assert len(results) == 1
        r = results[0]
        assert r.check_type == "unique"
        assert r.content == "COUNT(col) - COUNT(DISTINCT(col))"
        assert r.params == {"equal_to": 0, "accept_none": True}

    def test_accepted_range_with_tolerance(self):
        op = self._make_operator({"col": {"distinct_check": {"equal_to": 10, "tolerance": 0.1}}})
        op.column_mapping["col"]["distinct_check"]["result"] = 10
        op.column_mapping["col"]["distinct_check"]["success"] = True
        results = op._build_check_results()
        assert len(results) == 1
        r = results[0]
        assert r.check_type == "accepted_range"
        assert r.expected == ">= 9.0, <= 11.0"
        assert r.actual == "10"
        assert r.content == "COUNT(DISTINCT(col))"
        assert r.params == {"equal_to": 10, "tolerance": 0.1, "accept_none": True}

    def test_accepted_range_geq_leq(self):
        op = self._make_operator({"col": {"min": {"geq_to": 1, "leq_to": 100}}})
        op.column_mapping["col"]["min"]["result"] = 50
        op.column_mapping["col"]["min"]["success"] = True
        results = op._build_check_results()
        assert len(results) == 1
        r = results[0]
        assert r.check_type == "accepted_range"
        assert r.expected == ">=1, <=100"
        assert r.params == {"geq_to": 1, "leq_to": 100, "accept_none": True}

    def test_multiple_checks_correct_names_and_order(self):
        op = self._make_operator(
            {
                "col_a": {
                    "null_check": {"equal_to": 0, "result": 0, "success": True},
                    "min": {"geq_to": 1, "result": 5, "success": True},
                },
                "col_b": {
                    "max": {"less_than": 100, "result": 50, "success": True},
                },
            }
        )
        results = op._build_check_results()
        assert len(results) == 3
        assert results[0].name == "col_a.null_check"
        assert results[1].name == "col_a.min"
        assert results[2].name == "col_b.max"

    def test_failing_check_recorded_correctly(self):
        op = self._make_operator({"col": {"null_check": {"equal_to": 0}}})
        op.column_mapping["col"]["null_check"]["result"] = 5
        op.column_mapping["col"]["null_check"]["success"] = False
        results = op._build_check_results()
        assert len(results) == 1
        assert results[0].success is False
        assert results[0].actual == "5"

    def test_partition_clause_in_params(self):
        op = self._make_operator({"col": {"null_check": {"equal_to": 0, "partition_clause": "year=2024"}}})
        op.partition_clause = "region=us"
        op.column_mapping["col"]["null_check"]["result"] = 0
        op.column_mapping["col"]["null_check"]["success"] = True
        results = op._build_check_results()
        assert len(results) == 1
        assert results[0].params == {
            "equal_to": 0,
            "accept_none": True,
            "partition_clause": "region=us",
            "check_partition_clause": "year=2024",
        }

    def test_build_check_results_failure_returns_empty_list(self):
        op = self._make_operator({"col": {"null_check": {"equal_to": 0, "result": 0, "success": True}}})
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results()
        assert results == []

    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch.object(SQLColumnCheckOperator, "get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _):
        mock_hook.return_value.get_records.return_value = [("col", "null_check", 0)]
        op = self._make_operator({"col": {"null_check": {"equal_to": 0}}})
        op.execute(MagicMock())
        assert op.check_results == []

    @mock.patch.object(SQLColumnCheckOperator, "get_db_hook")
    def test_execute_populates_check_results(self, mock_hook):
        # Records returned by the DB: (column, check_name, result)
        mock_hook.return_value.get_records.return_value = [("col", "null_check", 0)]
        op = self._make_operator({"col": {"null_check": {"equal_to": 0}}})
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "col.null_check"
        assert r.check_type == "not_null"
        assert r.success is True
        assert r.severity == "error"
        assert r.column == "col"
        assert r.table == "test_table"
        assert r.expected == "0"
        assert r.actual == "0"
        assert r.content == "SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)"
        assert r.description == "Column-level statistical check against the configured threshold"
        assert r.params == {"equal_to": 0, "accept_none": True}

    @pytest.mark.parametrize("operator", ["greater_than", "geq_to", "less_than", "leq_to", "equal_to"])
    def test_known_comparison_operators_produce_expected_string(self, operator):
        """Regression guard for existing operators. Does NOT detect new operators being added —
        when adding a new comparison operator, add a case to this parametrize list manually."""
        result = SQLColumnCheckOperator._format_column_expected({operator: 5})
        assert result != ""

    @pytest.mark.parametrize(
        ("check", "check_values", "expected_type"),
        [
            # Range operators always win
            ("null_check", {"greater_than": 0}, "accepted_range"),
            ("col_count", {"geq_to": 10}, "accepted_range"),
            ("max", {"less_than": 100}, "accepted_range"),
            ("min", {"leq_to": 5}, "accepted_range"),
            # equal_to + non-zero tolerance → range (tolerance expands to bounds)
            ("distinct_check", {"equal_to": 10, "tolerance": 0.1}, "accepted_range"),
            # equal_to without tolerance → semantic names for standard zero-value checks
            ("null_check", {"equal_to": 0}, "not_null"),
            ("unique_check", {"equal_to": 0}, "unique"),
            # equal_to=0 with tolerance: condition requires equal_to != 0, so falls through
            ("null_check", {"equal_to": 0, "tolerance": 0.1}, "not_null"),
            ("distinct_check", {"equal_to": 0, "tolerance": 0.1}, "accepted_values"),
            # equal_to with non-zero value on standard check names
            ("null_check", {"equal_to": 5}, "accepted_values"),
            ("unique_check", {"equal_to": 5}, "accepted_values"),
            # generic check name with no comparison operators → returned unchanged
            ("custom_metric", {}, "custom_metric"),
            # equal_to on a non-null/unique check name
            ("distinct_check", {"equal_to": 10}, "accepted_values"),
        ],
    )
    def test_get_column_check_type(self, check, check_values, expected_type):
        assert SQLColumnCheckOperator._get_column_check_type(check, check_values) == expected_type

    @pytest.mark.parametrize(
        ("check_values", "expected"),
        [
            ({"equal_to": 0}, "0"),
            ({"equal_to": 5}, "5"),
            ({"greater_than": 5}, ">5"),
            ({"geq_to": 5}, ">=5"),
            ({"less_than": 10}, "<10"),
            ({"leq_to": 10}, "<=10"),
            ({"greater_than": 0, "less_than": 100}, ">0, <100"),
            ({"geq_to": 5, "leq_to": 10}, ">=5, <=10"),
        ],
    )
    def test_format_column_expected_no_tolerance(self, check_values, expected):
        assert SQLColumnCheckOperator._format_column_expected(check_values) == expected

    @pytest.mark.parametrize(
        ("check_values", "expected"),
        [
            ({"equal_to": 5, "tolerance": 0.1}, ">= 4.5, <= 5.5"),
            ({"equal_to": 10, "tolerance": 0.5}, ">= 5.0, <= 15.0"),
            ({"greater_than": 10, "tolerance": 0.1}, "> 9.0"),
            ({"geq_to": 10, "tolerance": 0.1}, ">= 9.0"),
            ({"less_than": 10, "tolerance": 0.1}, "< 11.0"),
            ({"leq_to": 10, "tolerance": 0.1}, "<= 11.0"),
            ({"geq_to": 10, "leq_to": 20, "tolerance": 0.1}, ">= 9.0, <= 22.0"),
        ],
    )
    def test_format_column_expected_with_tolerance(self, check_values, expected):
        assert SQLColumnCheckOperator._format_column_expected(check_values) == expected


class TestSQLTableCheckOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(checks, **kwargs):
        return SQLTableCheckOperator(task_id="test_task", table="test_table", checks=checks, **kwargs)

    def test_passing_check(self):
        checks = {"row_count_check": {"check_statement": "COUNT(*) >= 3"}}
        op = self._make_operator(checks)
        op.checks["row_count_check"]["result"] = "1"
        op.checks["row_count_check"]["success"] = True
        records = [("row_count_check", 1)]
        results = op._build_check_results(records)
        assert len(results) == 1
        r = results[0]
        assert r.name == "row_count_check"
        assert r.check_type == "expression_is_true"
        assert r.success is True
        assert r.table == "test_table"
        assert r.content == "COUNT(*) >= 3"
        assert r.expected == "all truthy"
        assert r.actual == "1"
        assert r.severity == "error"
        assert r.params == {"accept_none": False}

    def test_failing_check(self):
        checks = {"row_count_check": {"check_statement": "COUNT(*) >= 3"}}
        op = self._make_operator(checks)
        op.checks["row_count_check"]["result"] = "0"
        op.checks["row_count_check"]["success"] = False
        results = op._build_check_results([("row_count_check", 0)])
        assert len(results) == 1
        assert results[0].success is False
        assert results[0].actual == "0"

    def test_accept_none_severity_warn_when_empty_records(self):
        checks = {"row_count_check": {"check_statement": "COUNT(*) >= 3"}}
        op = self._make_operator(checks, accept_none=True)
        op.checks["row_count_check"]["success"] = False
        results = op._build_check_results([])
        assert len(results) == 1
        assert results[0].severity == "warn"

    def test_severity_error_when_records_present(self):
        checks = {"row_count_check": {"check_statement": "COUNT(*) >= 3"}}
        op = self._make_operator(checks, accept_none=True)
        op.checks["row_count_check"]["result"] = "1"
        op.checks["row_count_check"]["success"] = True
        results = op._build_check_results([("row_count_check", 1)])
        assert results[0].severity == "error"

    def test_multiple_checks(self):
        checks = {
            "count_check": {"check_statement": "COUNT(*) > 0"},
            "sum_check": {"check_statement": "SUM(val) > 0"},
        }
        op = self._make_operator(checks)
        op.checks["count_check"].update({"result": "1", "success": True})
        op.checks["sum_check"].update({"result": "1", "success": True})
        results = op._build_check_results([("count_check", 1), ("sum_check", 1)])
        assert len(results) == 2
        assert results[0].name == "count_check"
        assert results[1].name == "sum_check"

    def test_build_check_results_failure_returns_empty_list(self):
        checks = {"row_count_check": {"check_statement": "COUNT(*) >= 3"}}
        op = self._make_operator(checks)
        op.checks["row_count_check"].update({"result": "1", "success": True})
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results([("row_count_check", 1)])
        assert results == []

    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch.object(SQLTableCheckOperator, "get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _):
        mock_hook.return_value.get_records.return_value = [("row_count_check", 1)]
        op = self._make_operator({"row_count_check": {"check_statement": "COUNT(*) >= 1"}})
        op.execute(MagicMock())
        assert op.check_results == []

    @mock.patch.object(SQLTableCheckOperator, "get_db_hook")
    def test_execute_populates_check_results(self, mock_hook):
        mock_hook.return_value.get_records.return_value = [("row_count_check", "1")]
        op = self._make_operator({"row_count_check": {"check_statement": "COUNT(*) >= 1"}})
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "row_count_check"
        assert r.check_type == "expression_is_true"
        assert r.success is True
        assert r.severity == "error"
        assert r.column is None
        assert r.table == "test_table"
        assert r.expected == "all truthy"
        assert r.actual == "1"
        assert r.content == "COUNT(*) >= 1"
        assert r.description == "User-defined SQL expression must evaluate to true"
        assert r.params == {"accept_none": False}


class TestSQLCheckOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(**kwargs):
        return SQLCheckOperator(task_id="test_task", sql="SELECT 1", **kwargs)

    def test_all_truthy_records(self):
        op = self._make_operator()
        results = op._build_check_results([1, 1, 1])
        assert len(results) == 1
        r = results[0]
        assert r.name == "test_task"
        assert r.check_type == "expression_is_true"
        assert r.success is True
        assert r.expected == "all truthy"
        assert r.actual == "[1, 1, 1]"
        assert r.content == "SELECT 1"
        assert r.params is None

    def test_falsy_value_in_records(self):
        op = self._make_operator()
        results = op._build_check_results([1, 0, 1])
        assert len(results) == 1
        assert results[0].success is False

    def test_empty_records(self):
        op = self._make_operator()
        results = op._build_check_results([])
        assert len(results) == 1
        r = results[0]
        assert r.success is False
        assert r.actual is None

    def test_dict_records_all_true(self):
        op = self._make_operator()
        results = op._build_check_results({"A": True, "B": True})
        assert len(results) == 1
        assert results[0].success is True

    def test_dict_records_not_all_true(self):
        op = self._make_operator()
        results = op._build_check_results({"A": True, "B": False})
        assert len(results) == 1
        assert results[0].success is False

    def test_parameters_in_params(self):
        op = self._make_operator(parameters="my_params")
        results = op._build_check_results([1])
        assert len(results) == 1
        assert results[0].params == {"parameters": "my_params"}

    def test_build_check_results_failure_returns_empty_list(self):
        op = self._make_operator()
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results([1])
        assert results == []

    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch.object(SQLCheckOperator, "get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _):
        mock_hook.return_value.get_first.return_value = [1]
        op = self._make_operator()
        op.execute(MagicMock())
        assert op.check_results == []

    @mock.patch.object(SQLCheckOperator, "get_db_hook")
    def test_execute_populates_check_results(self, mock_hook):
        mock_hook.return_value.get_first.return_value = [1, 2, 3]
        op = self._make_operator()
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "test_task"
        assert r.check_type == "expression_is_true"
        assert r.success is True
        assert r.severity == "error"
        assert r.column is None
        assert r.table is None
        assert r.expected == "all truthy"
        assert r.actual == "[1, 2, 3]"
        assert r.content == "SELECT 1"
        assert r.description == "All values in the first returned row must evaluate to true"
        assert r.params is None


class TestSQLValueCheckOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(pass_value, tolerance=None):
        return SQLValueCheckOperator(
            task_id="test_task", sql="SELECT val FROM t", pass_value=pass_value, tolerance=tolerance
        )

    def test_exact_match_no_tolerance(self):
        op = self._make_operator(pass_value="5")
        results = op._build_check_results([5])
        assert len(results) == 1
        r = results[0]
        assert r.name == "test_task"
        assert r.check_type == "accepted_values"
        assert r.success is True
        assert r.expected == "5"
        assert r.actual == "[5]"
        assert r.content == "SELECT val FROM t"
        assert r.params == {"pass_value": "5"}

    def test_numeric_tolerance_produces_accepted_range(self):
        op = self._make_operator(pass_value=5, tolerance=0.1)
        results = op._build_check_results([5])
        assert len(results) == 1
        r = results[0]
        assert r.check_type == "accepted_range"
        assert r.expected == ">= 4.5, <= 5.5"
        assert r.params == {"pass_value": "5", "tolerance": 0.1}

    def test_non_numeric_pass_value_is_accepted_values(self):
        op = self._make_operator(pass_value="hello")
        results = op._build_check_results(["hello"])
        assert len(results) == 1
        assert results[0].check_type == "accepted_values"
        assert results[0].expected == "hello"

    def test_failing_check(self):
        op = self._make_operator(pass_value="10")
        results = op._build_check_results([99])
        assert len(results) == 1
        assert results[0].success is False

    def test_parameters_in_params(self):
        op = self._make_operator(pass_value="5")
        op.parameters = {"key": "val"}
        results = op._build_check_results([5])
        assert len(results) == 1
        assert results[0].params == {"pass_value": "5", "parameters": {"key": "val"}}

    def test_build_check_results_failure_returns_empty_list(self):
        op = self._make_operator(pass_value="5")
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results([5])
        assert results == []

    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch.object(SQLValueCheckOperator, "get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _):
        mock_hook.return_value.get_first.return_value = [5]
        op = self._make_operator(pass_value="5")
        op.execute(MagicMock())
        assert op.check_results == []

    @mock.patch.object(SQLValueCheckOperator, "get_db_hook")
    def test_execute_populates_check_results(self, mock_hook):
        mock_hook.return_value.get_first.return_value = [5]
        op = self._make_operator(pass_value="5")
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "test_task"
        assert r.check_type == "accepted_values"
        assert r.success is True
        assert r.severity == "error"
        assert r.column is None
        assert r.table is None
        assert r.expected == "5"
        assert r.actual == "[5]"
        assert r.content == "SELECT val FROM t"
        assert r.description == "All values in the first returned row must match the expected value"
        assert r.params == {"pass_value": "5"}

    @mock.patch.object(SQLValueCheckOperator, "get_db_hook")
    def test_execute_populates_check_results_with_tolerance(self, mock_hook):
        mock_hook.return_value.get_first.return_value = [5]
        op = self._make_operator(pass_value=5, tolerance=0.1)
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "test_task"
        assert r.check_type == "accepted_range"
        assert r.success is True
        assert r.severity == "error"
        assert r.column is None
        assert r.table is None
        assert r.expected == ">= 4.5, <= 5.5"
        assert r.actual == "[5]"
        assert r.content == "SELECT val FROM t"
        assert r.description == "All values in the first returned row must match the expected value"
        assert r.params == {"pass_value": "5", "tolerance": 0.1}


class TestSQLIntervalCheckOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(metrics_thresholds, **kwargs):
        return SQLIntervalCheckOperator(
            task_id="test_task",
            table="test_table",
            metrics_thresholds=metrics_thresholds,
            ratio_formula="max_over_min",
            ignore_zero=True,
            **kwargs,
        )

    def _make_all_tests_results(self, metric, current, past, threshold, ratio, success):
        return {
            metric: {
                "metric": metric,
                "current_metric": current,
                "past_metric": past,
                "threshold": threshold,
                "ignore_zero": True,
                "ratio": ratio,
                "success": success,
            }
        }

    def test_passing_metric(self):
        op = self._make_operator({"f1": 1.5})
        all_tests_results = self._make_all_tests_results("f1", 10, 9, 1.5, 1.1, True)
        results = op._build_check_results(all_tests_results)
        assert len(results) == 1
        r = results[0]
        assert r.name == "interval_f1"
        assert r.check_type == "accepted_range"
        assert r.success is True
        assert r.table == "test_table"
        assert r.expected == "< 1.5"
        assert r.actual == "1.1"
        assert r.content == "max(10, 9) / min(10, 9)"
        assert r.description == "Ratio of current metric to historical baseline must be below the threshold"
        assert r.params == {
            "threshold": 1.5,
            "days_back": -7,
            "ratio_formula_name": "max_over_min",
            "ratio_formula": "max({current}, {past}) / min({current}, {past})",
            "date_filter_column": "ds",
            "ignore_zero": True,
            "current_metric": 10,
            "past_metric": 9,
        }

    def test_failing_metric(self):
        op = self._make_operator({"f1": 1.5})
        all_tests_results = self._make_all_tests_results("f1", 10, 2, 1.5, 5.0, False)
        results = op._build_check_results(all_tests_results)
        assert len(results) == 1
        assert results[0].success is False
        assert results[0].actual == "5.0"

    def test_zero_ratio_none(self):
        op = self._make_operator({"f1": 1.5})
        all_tests_results = self._make_all_tests_results("f1", 0, 10, 1.5, None, True)
        results = op._build_check_results(all_tests_results)
        assert len(results) == 1
        assert results[0].actual == "0"
        assert results[0].success is True

    def test_multiple_metrics(self):
        op = self._make_operator({"f1": 1.5, "f2": 2.0})
        all_tests_results = {
            "f1": {
                "metric": "f1",
                "current_metric": 10,
                "past_metric": 9,
                "threshold": 1.5,
                "ignore_zero": True,
                "ratio": 1.1,
                "success": True,
            },
            "f2": {
                "metric": "f2",
                "current_metric": 10,
                "past_metric": 2,
                "threshold": 2.0,
                "ignore_zero": True,
                "ratio": 5.0,
                "success": False,
            },
        }
        results = op._build_check_results(all_tests_results)
        assert len(results) == 2
        assert results[0].name == "interval_f1"
        assert results[0].success is True
        assert results[1].name == "interval_f2"
        assert results[1].success is False

    def test_build_check_results_failure_returns_empty_list(self):
        op = self._make_operator({"f1": 1.5})
        all_tests_results = self._make_all_tests_results("f1", 10, 9, 1.5, 1.1, True)
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results(all_tests_results)
        assert results == []

    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch.object(SQLIntervalCheckOperator, "get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _):
        mock_hook.return_value.get_first.side_effect = [[10], [10]]
        op = self._make_operator({"f1": 1.5})
        op.execute(MagicMock())
        assert op.check_results == []

    @mock.patch.object(SQLIntervalCheckOperator, "get_db_hook")
    def test_execute_populates_check_results(self, mock_hook):
        # execute() fetches sql2 (past/reference) first, then sql1 (current)
        mock_hook.return_value.get_first.side_effect = [[9], [10]]
        op = self._make_operator({"f1": 1.5})
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "interval_f1"
        assert r.check_type == "accepted_range"
        assert r.success is True
        assert r.severity == "error"
        assert r.column is None
        assert r.table == "test_table"
        assert r.expected == "< 1.5"
        # ratio = max(10, 9) / min(10, 9) = 10/9 ≈ 1.111, content substitutes actual values
        assert r.content == "max(10, 9) / min(10, 9)"
        assert r.actual == str(10 / 9)
        assert r.description == "Ratio of current metric to historical baseline must be below the threshold"
        assert r.params == {
            "threshold": 1.5,
            "days_back": -7,
            "ratio_formula_name": "max_over_min",
            "ratio_formula": "max({current}, {past}) / min({current}, {past})",
            "date_filter_column": "ds",
            "ignore_zero": True,
            "current_metric": 10,
            "past_metric": 9,
        }

    def test_all_ratio_formulas_have_content_expressions(self):
        """Fails if a new ratio formula is added to ratio_formulas without a matching expression entry."""
        op = self._make_operator({"f1": 1.5})
        assert set(op.ratio_formulas.keys()) == set(op.ratio_formula_expressions.keys())


class TestSQLThresholdCheckOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(min_threshold=1, max_threshold=100):
        return SQLThresholdCheckOperator(
            task_id="test_task",
            sql="SELECT val FROM t",
            min_threshold=min_threshold,
            max_threshold=max_threshold,
        )

    def test_within_threshold(self):
        op = self._make_operator(min_threshold=1, max_threshold=100)
        meta_data = {"within_threshold": True, "min_threshold": 1.0, "max_threshold": 100.0}
        results = op._build_check_results(50, meta_data)
        assert len(results) == 1
        r = results[0]
        assert r.name == "test_task"
        assert r.check_type == "accepted_range"
        assert r.success is True
        assert r.expected == ">= 1.0, <= 100.0"
        assert r.actual == "50"
        assert r.content == "SELECT val FROM t"
        assert r.description == "SQL result must fall within the configured bounds"
        assert r.params == {"min_threshold": "1", "max_threshold": "100"}

    def test_outside_threshold(self):
        op = self._make_operator(min_threshold=20, max_threshold=100)
        meta_data = {"within_threshold": False, "min_threshold": 20.0, "max_threshold": 100.0}
        results = op._build_check_results(10, meta_data)
        assert len(results) == 1
        assert results[0].success is False
        assert results[0].actual == "10"
        assert results[0].expected == ">= 20.0, <= 100.0"

    def test_sql_threshold_raw_strings_preserved_in_params(self):
        op = self._make_operator(min_threshold="SELECT MIN(val) FROM ref", max_threshold=100)
        meta_data = {"within_threshold": True, "min_threshold": 5.0, "max_threshold": 100.0}
        results = op._build_check_results(50, meta_data)
        assert len(results) == 1
        assert results[0].params == {
            "min_threshold": "SELECT MIN(val) FROM ref",
            "max_threshold": "100",
        }

    def test_build_check_results_failure_returns_empty_list(self):
        op = self._make_operator()
        meta_data = {"within_threshold": True, "min_threshold": 1.0, "max_threshold": 100.0}
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results(50, meta_data)
        assert results == []

    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _):
        mock_hook.return_value.get_first.return_value = (50,)
        op = self._make_operator(min_threshold=1, max_threshold=100)
        op.execute(MagicMock())
        assert op.check_results == []

    @mock.patch.object(SQLThresholdCheckOperator, "get_db_hook")
    def test_execute_populates_check_results(self, mock_hook):
        mock_hook.return_value.get_first.return_value = (50,)
        op = self._make_operator(min_threshold=1, max_threshold=100)
        op.execute(MagicMock())
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "test_task"
        assert r.check_type == "accepted_range"
        assert r.success is True
        assert r.severity == "error"
        assert r.column is None
        assert r.table is None
        assert r.expected == ">= 1.0, <= 100.0"
        assert r.actual == "50"
        assert r.content == "SELECT val FROM t"
        assert r.description == "SQL result must fall within the configured bounds"
        assert r.params == {"min_threshold": "1", "max_threshold": "100"}


class TestBranchSQLOperatorBuildCheckResults:
    @staticmethod
    def _make_operator(**kwargs):
        return BranchSQLOperator(
            task_id="test_task",
            sql="SELECT 1",
            follow_task_ids_if_true=["branch_true"],
            follow_task_ids_if_false=["branch_false"],
            **kwargs,
        )

    def test_true_branch(self):
        op = self._make_operator()
        op.follow_branch = ["branch_true"]
        results = op._build_check_results(1)
        assert len(results) == 1
        r = results[0]
        assert r.name == "test_task"
        assert r.check_type == "expression_is_true"
        assert r.success is True
        assert r.severity == "info"
        assert r.expected == "truthy"
        assert r.actual == "1"
        assert r.content == "SELECT 1"
        assert r.description == "SQL result is evaluated as boolean to determine the execution branch"
        assert r.params == {
            "follow_task_ids_if_true": ["branch_true"],
            "follow_task_ids_if_false": ["branch_false"],
            "follow_branch": ["branch_true"],
        }

    def test_false_branch(self):
        op = self._make_operator()
        op.follow_branch = ["branch_false"]
        results = op._build_check_results(0)
        assert len(results) == 1
        r = results[0]
        assert r.success is False
        assert r.actual == "0"
        assert r.params == {
            "follow_task_ids_if_true": ["branch_true"],
            "follow_task_ids_if_false": ["branch_false"],
            "follow_branch": ["branch_false"],
        }

    def test_parameters_included_when_set(self):
        op = self._make_operator(parameters={"key": "val"})
        op.follow_branch = ["branch_true"]
        results = op._build_check_results(1)
        assert len(results) == 1
        assert results[0].params == {
            "follow_task_ids_if_true": ["branch_true"],
            "follow_task_ids_if_false": ["branch_false"],
            "follow_branch": ["branch_true"],
            "parameters": {"key": "val"},
        }

    def test_build_check_results_failure_returns_empty_list(self):
        op = self._make_operator()
        op.follow_branch = ["branch_true"]
        with mock.patch(
            "airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom")
        ):
            results = op._build_check_results(1)
        assert results == []

    @mock.patch.object(BranchSQLOperator, "skip_all_except")
    @mock.patch("airflow.providers.common.sql.operators.sql.SQLCheckResult", side_effect=RuntimeError("boom"))
    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_execute_unaffected_when_build_check_results_raises(self, mock_hook, _, mock_skip):
        mock_hook.return_value.get_first.return_value = 1
        op = self._make_operator()
        op.execute({"ti": MagicMock()})
        assert op.check_results == []
        mock_skip.assert_called_once()

    @mock.patch.object(BranchSQLOperator, "skip_all_except")
    @mock.patch("airflow.providers.common.sql.operators.sql.BaseSQLOperator.get_db_hook")
    def test_execute_populates_check_results(self, mock_hook, mock_skip):
        mock_hook.return_value.get_first.return_value = 1
        op = self._make_operator()
        op.execute({"ti": MagicMock()})
        assert len(op.check_results) == 1
        r = op.check_results[0]
        assert r.name == "test_task"
        assert r.check_type == "expression_is_true"
        assert r.success is True
        assert r.severity == "info"
        assert r.column is None
        assert r.table is None
        assert r.expected == "truthy"
        assert r.actual == "1"
        assert r.content == "SELECT 1"
        assert r.description == "SQL result is evaluated as boolean to determine the execution branch"
        assert r.params == {
            "follow_task_ids_if_true": ["branch_true"],
            "follow_task_ids_if_false": ["branch_false"],
            "follow_branch": ["branch_true"],
        }
        mock_skip.assert_called_once()


class TestSqlBaseOperatorAttachCheckFacets:
    """Tests for BaseSQLOperator._attach_check_facets."""

    @staticmethod
    def _make_operator():
        return SQLCheckOperator(task_id="test_task", sql="SELECT 1")

    @staticmethod
    def _dataset(name):
        from openlineage.client.event_v2 import Dataset

        return Dataset(namespace="default", name=name)

    def test_empty_check_results_returns_lineage_unchanged(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        result = op._attach_check_facets(OperatorLineage())
        assert result == OperatorLineage()
        assert result.run_facets == {}

    def test_old_openlineage_client_raises_optional_feature_exception(self):
        try:
            from airflow.providers.openlineage.extractors import OperatorLineage
        except ImportError:
            pytest.skip("openlineage provider not installed")
        from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException

        op = self._make_operator()
        op.check_results = [SQLCheckResult(name="t", check_type="expr", success=True)]
        with mock.patch(
            "airflow.providers.common.compat.openlineage.check.metadata.version", return_value="1.46.0"
        ):
            with pytest.raises(AirflowOptionalProviderFeatureException):
                op._attach_check_facets(OperatorLineage())

    def test_run_facet_all_fields_populated(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="my_check",
                check_type="expression_is_true",
                success=True,
                severity="warn",
                column="col_a",
                table=None,  # no table → run facet
                expected="all truthy",
                actual="[1, 1]",
                content="SELECT col_a FROM t",
                description="All values must be truthy",
                params={"custom_key": "custom_val"},
            )
        ]
        result = op._attach_check_facets(OperatorLineage())

        facet = result.run_facets.get("test")
        assert facet is not None
        assert len(facet.tests) == 1
        t = facet.tests[0]
        assert t.name == "my_check"
        assert t.status == "pass"
        assert t.severity == "warn"
        assert t.type == "expression_is_true"
        assert t.description == "All values must be truthy"
        assert t.expected == "all truthy"
        assert t.actual == "[1, 1]"
        assert t.content == "SELECT col_a FROM t"
        assert t.contentType == "sql"
        # params merges tested_column / tested_table with the SQLCheckResult params
        assert t.params == {
            "tested_column": "col_a",
            "tested_table": None,
            "custom_key": "custom_val",
        }

    def test_run_facet_failing_check_has_fail_status(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [SQLCheckResult(name="bad_check", check_type="not_null", success=False)]
        result = op._attach_check_facets(OperatorLineage())

        facet = result.run_facets.get("test")
        assert facet is not None
        assert facet.tests[0].status == "fail"

    def test_dataset_facet_all_fields_populated(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="col.null_check",
                check_type="not_null",
                success=True,
                severity="error",
                column="col",
                table="test_table",
                expected="0 nulls",
                actual="0",
                content="SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)",
                description="Column must contain no nulls",
                params={"accept_none": False},
            )
        ]
        lineage = OperatorLineage(inputs=[self._dataset("myschema.test_table")])
        result = op._attach_check_facets(lineage)

        assert result.run_facets == {}
        dq_facet = result.inputs[0].facets.get("dataQualityAssertions")
        assert dq_facet is not None
        assert len(dq_facet.assertions) == 1
        a = dq_facet.assertions[0]
        assert a.assertion == "not_null"
        assert a.success is True
        assert a.severity == "error"
        assert a.column == "col"
        assert a.name == "col.null_check"
        assert a.description == "Column must contain no nulls"
        assert a.expected == "0 nulls"
        assert a.actual == "0"
        assert a.content == "SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)"
        assert a.contentType == "sql"
        assert a.params == {"accept_none": False}

    def test_dataset_facet_multiple_assertions_same_table(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="col_a.null_check",
                check_type="not_null",
                success=True,
                column="col_a",
                table="test_table",
                expected="0 nulls",
                actual="0",
            ),
            SQLCheckResult(
                name="col_b.unique_check",
                check_type="unique",
                success=False,
                column="col_b",
                table="test_table",
                expected="0 duplicates",
                actual="3",
            ),
            SQLCheckResult(
                name="row_count_check",
                check_type="expression_is_true",
                success=True,
                column=None,
                table="test_table",
                expected="all truthy",
                actual="1",
                content="COUNT(*) > 0",
            ),
        ]
        lineage = OperatorLineage(inputs=[self._dataset("myschema.test_table")])
        result = op._attach_check_facets(lineage)

        assert result.run_facets == {}
        dq_facet = result.inputs[0].facets.get("dataQualityAssertions")
        assert dq_facet is not None
        assert len(dq_facet.assertions) == 3

        a0 = dq_facet.assertions[0]
        assert a0.name == "col_a.null_check"
        assert a0.assertion == "not_null"
        assert a0.success is True
        assert a0.column == "col_a"
        assert a0.expected == "0 nulls"
        assert a0.actual == "0"

        a1 = dq_facet.assertions[1]
        assert a1.name == "col_b.unique_check"
        assert a1.assertion == "unique"
        assert a1.success is False
        assert a1.column == "col_b"
        assert a1.expected == "0 duplicates"
        assert a1.actual == "3"

        a2 = dq_facet.assertions[2]
        assert a2.name == "row_count_check"
        assert a2.assertion == "expression_is_true"
        assert a2.success is True
        assert a2.column is None
        assert a2.content == "COUNT(*) > 0"

    def test_dataset_facet_attached_to_output_all_fields_populated(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="col.null_check",
                check_type="not_null",
                success=True,
                severity="error",
                column="col",
                table="target_table",
                expected="0 nulls",
                actual="0",
                content="SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)",
                description="Column must contain no nulls",
                params={"accept_none": False},
            )
        ]
        # table appears only in outputs (e.g. a post-write quality check)
        lineage = OperatorLineage(outputs=[self._dataset("myschema.target_table")])
        result = op._attach_check_facets(lineage)

        assert result.run_facets == {}
        dq_facet = result.outputs[0].facets.get("dataQualityAssertions")
        assert dq_facet is not None
        assert len(dq_facet.assertions) == 1
        a = dq_facet.assertions[0]
        assert a.assertion == "not_null"
        assert a.success is True
        assert a.severity == "error"
        assert a.column == "col"
        assert a.name == "col.null_check"
        assert a.description == "Column must contain no nulls"
        assert a.expected == "0 nulls"
        assert a.actual == "0"
        assert a.content == "SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)"
        assert a.contentType == "sql"
        assert a.params == {"accept_none": False}

    def test_unmatched_table_falls_back_to_run_facet_with_table_context_in_params(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="col.null_check",
                check_type="not_null",
                success=False,
                severity="error",
                column="col",
                table="other_table",
                expected="0 nulls",
                actual="5",
                content="SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)",
                description="Column must contain no nulls",
                params={"accept_none": False},
            )
        ]
        lineage = OperatorLineage(inputs=[self._dataset("myschema.test_table")])
        result = op._attach_check_facets(lineage)

        assert result.inputs[0].facets == {}
        facet = result.run_facets.get("test")
        assert facet is not None
        assert len(facet.tests) == 1
        t = facet.tests[0]
        assert t.name == "col.null_check"
        assert t.status == "fail"
        assert t.severity == "error"
        assert t.type == "not_null"
        assert t.description == "Column must contain no nulls"
        assert t.expected == "0 nulls"
        assert t.actual == "5"
        assert t.content == "SUM(CASE WHEN col IS NULL THEN 1 ELSE 0 END)"
        assert t.contentType == "sql"
        # tested_column / tested_table are injected alongside the SQLCheckResult params
        assert t.params == {
            "tested_column": "col",
            "tested_table": "other_table",
            "accept_none": False,
        }

    def test_exact_match_preferred_over_suffix_match(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="col.null_check",
                check_type="not_null",
                success=True,
                column="col",
                table="orders",
            )
        ]
        suffix_match = self._dataset("schema.orders")  # endswith match — listed first
        exact_match = self._dataset("orders")  # exact match — listed second
        lineage = OperatorLineage(inputs=[suffix_match, exact_match])
        result = op._attach_check_facets(lineage)

        assert result.run_facets == {}
        # confirm ordering is preserved so index assertions below are unambiguous
        assert result.inputs[0].name == "schema.orders"
        assert result.inputs[1].name == "orders"
        # suffix-match dataset must not receive the facet
        assert result.inputs[0].facets == {}
        # exact-match dataset must receive it
        dq_facet = result.inputs[1].facets.get("dataQualityAssertions")
        assert dq_facet is not None
        assert dq_facet.assertions[0].name == "col.null_check"

    def test_same_table_in_inputs_and_outputs_both_receive_assertions(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            SQLCheckResult(
                name="col.null_check",
                check_type="not_null",
                success=True,
                column="col",
                table="orders",
                expected="0 nulls",
                actual="0",
            )
        ]
        input_ds = self._dataset("schema.orders")
        output_ds = self._dataset("schema.orders")
        lineage = OperatorLineage(inputs=[input_ds], outputs=[output_ds])
        result = op._attach_check_facets(lineage)

        assert result.run_facets == {}
        assert len(result.inputs) == 1
        assert len(result.outputs) == 1

        for ds in (result.inputs[0], result.outputs[0]):
            dq_facet = ds.facets.get("dataQualityAssertions")
            assert dq_facet is not None, f"missing facet on {ds.name}"
            assert len(dq_facet.assertions) == 1
            a = dq_facet.assertions[0]
            assert a.name == "col.null_check"
            assert a.assertion == "not_null"
            assert a.success is True
            assert a.column == "col"

    def test_mixed_results_dataset_facet_and_run_facet_populated_correctly(self):
        pytest.importorskip(
            "openlineage.client", minversion="1.47.0", reason="openlineage-python >= 1.47.0 required"
        )
        from airflow.providers.openlineage.extractors import OperatorLineage

        op = self._make_operator()
        op.check_results = [
            # matched input dataset → DataQualityAssertionsDatasetFacet
            SQLCheckResult(
                name="col_a.null_check",
                check_type="not_null",
                success=True,
                column="col_a",
                table="known_table",
                expected="0 nulls",
                actual="0",
            ),
            # table not in any dataset → falls back to TestRunFacet
            SQLCheckResult(
                name="col_b.null_check",
                check_type="not_null",
                success=False,
                column="col_b",
                table="unknown_table",
                expected="0 nulls",
                actual="2",
            ),
            # no table at all → directly to TestRunFacet
            SQLCheckResult(
                name="row_count_check",
                check_type="expression_is_true",
                success=True,
                table=None,
                expected="all truthy",
                actual="1",
                content="COUNT(*) > 0",
            ),
        ]
        lineage = OperatorLineage(inputs=[self._dataset("myschema.known_table")])
        result = op._attach_check_facets(lineage)

        # matched check lands on the input dataset facet
        dq_facet = result.inputs[0].facets.get("dataQualityAssertions")
        assert dq_facet is not None
        assert len(dq_facet.assertions) == 1
        a = dq_facet.assertions[0]
        assert a.name == "col_a.null_check"
        assert a.assertion == "not_null"
        assert a.success is True
        assert a.column == "col_a"

        # unmatched-table check and no-table check both land in the run facet
        run_facet = result.run_facets.get("test")
        assert run_facet is not None
        assert len(run_facet.tests) == 2

        t0 = run_facet.tests[0]
        assert t0.name == "row_count_check"
        assert t0.status == "pass"
        assert t0.params["tested_column"] is None
        assert t0.params["tested_table"] is None

        t1 = run_facet.tests[1]
        assert t1.name == "col_b.null_check"
        assert t1.status == "fail"
        assert t1.params["tested_column"] == "col_b"
        assert t1.params["tested_table"] == "unknown_table"


class TestSQLInsertRowsOperator:
    @mock.patch.object(SQLInsertRowsOperator, "get_db_hook")
    def test_insert_rows_operator_with_preoperator(self, mock_get_db_hook):
        mock_hook = MagicMock(spec=DbApiHook)
        mock_get_db_hook.return_value = mock_hook

        operator = SQLInsertRowsOperator(
            task_id="test_task",
            conn_id="default_conn",
            schema="hollywood",
            table_name="actors",
            preoperator="TRUNCATE TABLE hollywood.actors",
            rows=[
                (1, "Stallone", "Sylvester", 78),
                (2, "Statham", "Jason", 57),
                (3, "Li", "Jet", 61),
                (4, "Lundgren", "Dolph", 66),
                (5, "Norris", "Chuck", 84),
            ],
        )

        operator.execute({})

        mock_hook.run.assert_called_once()
        args, _ = mock_hook.run.call_args
        assert args[0] == "TRUNCATE TABLE hollywood.actors"

        mock_hook.insert_rows.assert_called_once()
        _, kwargs = mock_hook.insert_rows.call_args

        assert kwargs["rows"] == [
            (1, "Stallone", "Sylvester", 78),
            (2, "Statham", "Jason", 57),
            (3, "Li", "Jet", 61),
            (4, "Lundgren", "Dolph", 66),
            (5, "Norris", "Chuck", 84),
        ]

    @mock.patch.object(SQLInsertRowsOperator, "get_db_hook")
    def test_insert_rows_operator_with_rows_processor(self, mock_get_db_hook):
        mock_hook = MagicMock(spec=DbApiHook)
        mock_get_db_hook.return_value = mock_hook

        operator = SQLInsertRowsOperator(
            task_id="test_task",
            conn_id="default_conn",
            schema="hollywood",
            table_name="actors",
            rows=iter(
                [
                    {"index": 1, "name": "Stallone", "firstname": "Sylvester", "age": 78},
                    {"index": 2, "name": "Statham", "firstname": "Jason", "age": 57},
                    {"index": 3, "name": "Li", "firstname": "Jet", "age": 61},
                    {"index": 4, "name": "Lundgren", "firstname": "Dolph", "age": 66},
                    {"index": 5, "name": "Norris", "firstname": "Chuck", "age": 84},
                ]
            ),
            rows_processor=lambda rows, **context: [tuple(row.values()) for row in rows],
        )

        operator.execute({})

        mock_hook.insert_rows.assert_called_once()
        _, kwargs = mock_hook.insert_rows.call_args

        assert kwargs["rows"] == [
            (1, "Stallone", "Sylvester", 78),
            (2, "Statham", "Jason", 57),
            (3, "Li", "Jet", 61),
            (4, "Lundgren", "Dolph", 66),
            (5, "Norris", "Chuck", 84),
        ]
