test_dbt_parser.py

from pathlib import Path
import unittest
from unittest.mock import Mock, mock_open, patch

from data_diff.errors import (
    DataDiffDbtBigQueryUnsupportedMethodError,
    DataDiffDbtConnectionNotImplementedError,
    DataDiffDbtCoreNoRunnerError,
    DataDiffDbtNoSuccessfulModelsInRunError,
    DataDiffDbtProfileNotFoundError,
    DataDiffDbtRedshiftPasswordOnlyError,
    DataDiffDbtRunResultsVersionError,
    DataDiffDbtSnowflakeSetConnectionError,
)

from data_diff.dbt import (
    DbtParser,
)
from data_diff.dbt_parser import (
    RUN_RESULTS_PATH,
    PROJECT_FILE,
    TDatadiffConfig,
)


class TestDbtParser(unittest.TestCase):
    def test_get_datadiff_config(self):
        project_dict = {"vars": {"data_diff": {"prod_database": "a_prod_database"}}}

        mock_self = Mock()
        mock_self.project_dict = project_dict
        config = DbtParser.get_datadiff_config(mock_self)

        self.assertEqual(project_dict["vars"]["data_diff"]["prod_database"], config.prod_database)
        self.assertEqual(config.prod_schema, None)

    def test_get_datadiff_config_no_config(self):
        project_dict = {"key": {"key": "value"}}

        mock_self = Mock()
        mock_self.project_dict = project_dict

        config = DbtParser.get_datadiff_config(mock_self)
        self.assertEqual(config, TDatadiffConfig())

    def test_get_models(self):
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_self.dbt_version = "1.5.0"
        selection = "model+"
        mock_return_value = Mock()
        mock_self.get_dbt_selection_models.return_value = mock_return_value

        models = DbtParser.get_models(mock_self, selection)
        mock_self.get_dbt_selection_models.assert_called_once_with(selection)
        self.assertEqual(models, mock_return_value)

    def test_get_models_simple_select(self):
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_self.dbt_version = "1.4.0"
        selection = "model+"
        mock_return_value = Mock()
        mock_self.get_simple_model_selection.return_value = mock_return_value

        models = DbtParser.get_models(mock_self, selection)
        mock_self.get_dbt_selection_models.assert_not_called()
        mock_self.get_simple_model_selection.assert_called_with(selection)
        self.assertEqual(models, mock_return_value)

    def test_get_models_no_runner(self):
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_self.dbt_version = "1.5.0"
        mock_self.dbt_runner = None
        selection = "model+"
        mock_return_value = Mock()
        mock_self.get_dbt_selection_models.return_value = mock_return_value

        with self.assertRaises(DataDiffDbtCoreNoRunnerError):
            _ = DbtParser.get_models(mock_self, selection)
        mock_self.get_dbt_selection_models.assert_not_called()

    def test_get_models_no_selection(self):
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_self.dbt_version = "1.5.0"
        selection = None
        mock_return_value = Mock()
        mock_self.get_run_results_models.return_value = mock_return_value

        models = DbtParser.get_models(mock_self, selection)
        mock_self.get_dbt_selection_models.assert_not_called()
        mock_self.get_run_results_models.assert_called()
        self.assertEqual(models, mock_return_value)

    @patch("data_diff.dbt_parser.RunResultsJsonConfig.parse_obj")
    @patch("builtins.open", new_callable=mock_open, read_data="{}")
    def test_get_run_results_models(self, mock_open, mock_artifact_parser):
        mock_model = {"success_unique_id": "expected_value"}
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_run_results = Mock()
        mock_success_result = Mock()
        mock_success_result.status = mock_success_result.Status.success
        mock_fail_result = Mock()
        mock_fail_result.status = mock_fail_result.Status.fail
        mock_artifact_parser.return_value = mock_run_results
        mock_run_results.metadata.dbt_version = "1.0.0"
        mock_success_result.unique_id = "success_unique_id"
        mock_fail_result.unique_id = "fail_unique_id"
        mock_run_results.results = [mock_success_result, mock_fail_result]

        mock_self.dev_manifest_obj.nodes.get.return_value = mock_model
        models = DbtParser.get_run_results_models(mock_self)
        self.assertEqual(mock_model, models[0])
        mock_open.assert_called_with(Path(RUN_RESULTS_PATH))
        mock_artifact_parser.assert_called_once_with({})

    @patch("data_diff.dbt_parser.RunResultsJsonConfig.parse_obj")
    @patch("builtins.open", new_callable=mock_open, read_data="{}")
    def test_get_run_results_models_bad_lower_dbt_version(self, mock_open, mock_artifact_parser):
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_run_results = Mock()
        mock_artifact_parser.return_value = mock_run_results
        mock_run_results.metadata.dbt_version = "0.19.0"

        with self.assertRaises(DataDiffDbtRunResultsVersionError) as ex:
            DbtParser.get_run_results_models(mock_self)

        mock_open.assert_called_once_with(Path(RUN_RESULTS_PATH))
        mock_artifact_parser.assert_called_once_with({})
        self.assertIn("version to be", ex.exception.args[0])

    @patch("data_diff.dbt_parser.RunResultsJsonConfig.parse_obj")
    @patch("builtins.open", new_callable=mock_open, read_data="{}")
    def test_get_run_results_models_no_success(self, mock_open, mock_artifact_parser):
        mock_self = Mock()
        mock_self.project_dir = Path()
        mock_run_results = Mock()
        mock_fail_result = Mock()
        mock_fail_result.status = mock_fail_result.Status.fail
        mock_artifact_parser.return_value = mock_run_results
        mock_run_results.metadata.dbt_version = "1.0.0"
        mock_fail_result.unique_id = "fail_unique_id"
        mock_run_results.results = [mock_fail_result]

        with self.assertRaises(DataDiffDbtNoSuccessfulModelsInRunError):
            DbtParser.get_run_results_models(mock_self)

        mock_open.assert_any_call(Path(RUN_RESULTS_PATH))
        mock_artifact_parser.assert_called_once_with({})

    @patch("data_diff.dbt_parser.yaml")
    @patch("builtins.open", new_callable=mock_open, read_data="key:\n  value")
    def test_get_project_dict(self, mock_open, mock_yaml):
        expected_dict = {"key1": "value1"}
        mock_self = Mock()

        mock_self.project_dir = Path()
        mock_yaml.safe_load.return_value = expected_dict
        project_dict = DbtParser.get_project_dict(mock_self)

        self.assertEqual(project_dict, expected_dict)
        mock_open.assert_called_once_with(Path(PROJECT_FILE))

    def test_set_connection_snowflake_success_password(self):
        expected_driver = "snowflake"
        expected_credentials = {"user": "user", "password": "password"}
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        mock_self.set_casing_policy_for.assert_called_once_with("snowflake")
        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("user"), expected_credentials["user"])
        self.assertEqual(mock_self.connection.get("password"), expected_credentials["password"])
        self.assertEqual(mock_self.connection.get("key"), None)

    def test_set_connection_snowflake_success_key(self):
        expected_driver = "snowflake"
        expected_credentials = {"user": "user", "private_key_path": "private_key_path"}
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        mock_self.set_casing_policy_for.assert_called_once_with("snowflake")
        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("user"), expected_credentials["user"])
        self.assertEqual(mock_self.connection.get("password"), None)
        self.assertEqual(mock_self.connection.get("key"), expected_credentials["private_key_path"])

    def test_set_connection_snowflake_success_key_and_passphrase(self):
        expected_driver = "snowflake"
        expected_credentials = {
            "user": "user",
            "private_key_path": "private_key_path",
            "private_key_passphrase": "private_key_passphrase",
        }
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        mock_self.set_casing_policy_for.assert_called_once_with("snowflake")
        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("user"), expected_credentials["user"])
        self.assertEqual(mock_self.connection.get("password"), None)
        self.assertEqual(mock_self.connection.get("key"), expected_credentials["private_key_path"])
        self.assertEqual(
            mock_self.connection.get("private_key_passphrase"), expected_credentials["private_key_passphrase"]
        )

    def test_set_connection_snowflake_no_key_or_password(self):
        expected_driver = "snowflake"
        expected_credentials = {"user": "user"}
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        with self.assertRaises(DataDiffDbtSnowflakeSetConnectionError):
            DbtParser.set_connection(mock_self)

        self.assertNotIsInstance(mock_self.connection, dict)

    def test_set_connection_snowflake_authenticator(self):
        expected_driver = "snowflake"
        expected_credentials = {"user": "user", "authenticator": "authenticator"}
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("authenticator"), expected_credentials["authenticator"])
        self.assertEqual(mock_self.connection.get("user"), expected_credentials["user"])

    def test_set_connection_snowflake_key_and_password(self):
        expected_driver = "snowflake"
        expected_credentials = {"user": "user", "private_key_path": "private_key_path", "password": "password"}
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        with self.assertRaises(DataDiffDbtSnowflakeSetConnectionError):
            DbtParser.set_connection(mock_self)

        self.assertNotIsInstance(mock_self.connection, dict)

    def test_set_connection_bigquery_oauth(self):
        expected_driver = "bigquery"
        expected_credentials = {
            "method": "oauth",
            "project": "a_project",
            "dataset": "a_dataset",
        }
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("project"), expected_credentials["project"])
        self.assertEqual(mock_self.connection.get("dataset"), expected_credentials["dataset"])

    def test_set_connection_bigquery_oauth_sa_impersonation(self):
        expected_driver = "bigquery"
        expected_credentials = {
            "method": "oauth",
            "project": "a_project",
            "dataset": "a_dataset",
            "impersonate_service_account": "a_service_account@yourproject.iam.gserviceaccount.com",
        }
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("project"), expected_credentials["project"])
        self.assertEqual(mock_self.connection.get("dataset"), expected_credentials["dataset"])
        self.assertEqual(
            mock_self.connection.get("impersonate_service_account"),
            expected_credentials["impersonate_service_account"],
        )

    def test_set_connection_bigquery_svc_account(self):
        expected_driver = "bigquery"
        expected_credentials = {
            "method": "service-account",
            "project": "a_project",
            "dataset": "a_dataset",
            "keyfile": "/some/path",
        }
        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)

        DbtParser.set_connection(mock_self)

        self.assertIsInstance(mock_self.connection, dict)
        self.assertEqual(mock_self.connection.get("driver"), expected_driver)
        self.assertEqual(mock_self.connection.get("project"), expected_credentials["project"])
        self.assertEqual(mock_self.connection.get("dataset"), expected_credentials["dataset"])
        self.assertEqual(mock_self.connection.get("keyfile"), expected_credentials["keyfile"])

    def test_set_connection_bigquery_not_supported(self):
        expected_driver = "bigquery"
        expected_credentials = {
            "method": "not_supported",
            "project": "a_project",
            "dataset": "a_dataset",
        }

        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)
        with self.assertRaises(DataDiffDbtBigQueryUnsupportedMethodError):
            DbtParser.set_connection(mock_self)

        self.assertNotIsInstance(mock_self.connection, dict)

    def test_set_connection_redshift_not_password(self):
        driver = "redshift"
        credentials = {
            "method": "not_password",
        }

        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (credentials, driver)
        with self.assertRaises(DataDiffDbtRedshiftPasswordOnlyError):
            DbtParser.set_connection(mock_self)

        self.assertNotIsInstance(mock_self.connection, dict)

    def test_set_connection_not_implemented(self):
        expected_driver = "unimplemented_provider"

        mock_self = Mock()
        mock_self.get_connection_creds.return_value = (None, expected_driver)
        with self.assertRaises(DataDiffDbtConnectionNotImplementedError):
            DbtParser.set_connection(mock_self)

        self.assertNotIsInstance(mock_self.connection, dict)

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_creds_success(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {
            "a_profile": {
                "outputs": {
                    "a_target": {"type": "TYPE1", "credential_1": "credential_1", "credential_2": "credential_2"}
                },
                "target": "a_target",
            }
        }
        target = profiles_dict["a_profile"]["target"]
        expected_credentials = profiles_dict["a_profile"]["outputs"]["a_target"]
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        mock_self.project_dict = {"profile": "a_profile"}
        mock_yaml.safe_load.return_value = profiles_dict
        mock_profile_renderer().render_data.side_effect = [target, expected_credentials]
        credentials, conn_type = DbtParser.get_connection_creds(mock_self)
        self.assertEqual(credentials, expected_credentials)
        self.assertEqual(conn_type, "type1")

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_no_matching_profile(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {"a_profile": {}}
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        mock_self.project_dict = {"profile": "wrong_profile"}
        mock_yaml.safe_load.return_value = profiles_dict
        profile = profiles_dict["a_profile"]
        mock_profile_renderer().render_data.return_value = profile
        with self.assertRaises(DataDiffDbtProfileNotFoundError):
            _, _ = DbtParser.get_connection_creds(mock_self)

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_no_target(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {
            "a_profile": {
                "outputs": {
                    "a_target": {"type": "TYPE1", "credential_1": "credential_1", "credential_2": "credential_2"}
                },
            }
        }
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        profile = profiles_dict["a_profile"]
        mock_profile_renderer().render_data.return_value = profile
        mock_self.project_dict = {"profile": "a_profile"}
        mock_yaml.safe_load.return_value = profiles_dict
        with self.assertRaises(DataDiffDbtProfileNotFoundError):
            _, _ = DbtParser.get_connection_creds(mock_self)

    profile_yaml_no_outputs = """
    a_profile:
      target: a_target
    """

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_no_outputs(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {"a_profile": {"target": "a_target"}}
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        mock_self.project_dict = {"profile": "a_profile"}
        profile = profiles_dict["a_profile"]
        mock_profile_renderer().render_data.return_value = profile
        mock_yaml.safe_load.return_value = profiles_dict
        with self.assertRaises(DataDiffDbtProfileNotFoundError):
            _, _ = DbtParser.get_connection_creds(mock_self)

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_no_credentials(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {
            "a_profile": {
                "outputs": {"a_target": {}},
                "target": "a_target",
            }
        }
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        mock_self.project_dict = {"profile": "a_profile"}
        mock_yaml.safe_load.return_value = profiles_dict
        profile_target = profiles_dict["a_profile"]["target"]
        mock_profile_renderer().render_data.return_value = profile_target
        with self.assertRaises(DataDiffDbtProfileNotFoundError):
            _, _ = DbtParser.get_connection_creds(mock_self)

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_no_target_credentials(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {
            "a_profile": {
                "outputs": {
                    "a_target": {"type": "TYPE1", "credential_1": "credential_1", "credential_2": "credential_2"}
                },
                "target": "a_different_target",
            }
        }
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        mock_self.project_dict = {"profile": "a_profile"}
        profile_target = profiles_dict["a_profile"]["target"]
        mock_profile_renderer().render_data.return_value = profile_target
        mock_yaml.safe_load.return_value = profiles_dict
        with self.assertRaises(DataDiffDbtProfileNotFoundError):
            _, _ = DbtParser.get_connection_creds(mock_self)

    @patch("data_diff.dbt_parser.yaml")
    @patch("data_diff.dbt_parser.ProfileRenderer")
    @patch("builtins.open", new_callable=mock_open, read_data="")
    def test_get_connection_no_type(self, mock_open, mock_profile_renderer, mock_yaml):
        profiles_dict = {
            "a_profile": {
                "outputs": {"a_target": {"credential_1": "credential_1", "credential_2": "credential_2"}},
                "target": "a_target",
            }
        }
        mock_self = Mock()
        mock_self.profiles_dir = Path()
        mock_self.project_dict = {"profile": "a_profile"}
        mock_yaml.safe_load.return_value = profiles_dict
        profile_target = profiles_dict["a_profile"]["target"]
        mock_profile_renderer().render_data.return_value = profile_target
        with self.assertRaises(DataDiffDbtProfileNotFoundError):
            _, _ = DbtParser.get_connection_creds(mock_self)