custom exceptions and tests

Author: dlawin

data_diff/dbt.py

@@ -1,18 +1,17 @@
 import os
 import re
 import time
-import webbrowser
 from typing import List, Optional, Dict, Tuple, Union
 import keyring
 import pydantic
 import rich
-from rich.prompt import Confirm, Prompt
+from rich.prompt import Prompt
 
-from data_diff.errors import DataDiffCustomSchemaNoConfigError, DataDiffDbtProjectVarsNotFoundError
+from data_diff.errors import DataDiffCustomSchemaNoConfigError, DataDiffDbtProjectVarsNotFoundError, DataDiffNoAPIKeyError, DataDiffNoDatasourceIdError
 
 from . import connect_to_table, diff_tables, Algorithm
 from .cloud import DatafoldAPI, TCloudApiDataDiff, TCloudApiOrgMeta, get_or_create_data_source
-from .dbt_parser import DbtParser, PROJECT_FILE, TDatadiffConfig
+from .dbt_parser import DbtParser, TDatadiffConfig
 from .tracking import (
     bool_ask_for_email,
     create_email_signup_event_json,
@@ -38,6 +37,7 @@
 )
 
 logger = getLogger(__name__)
+CLOUD_DOC_URL = "https://docs.datafold.com/development_testing/cloud"
 
 
 class TDiffVars(pydantic.BaseModel):
@@ -81,11 +81,9 @@ def dbt_diff(
         org_meta = api.get_org_meta()
 
         if config.datasource_id is None:
-            cloud_doc_url = "https://docs.datafold.com/development_testing/cloud"
             rich.print("[red]Data source ID not found in dbt_project.yml")
-            raise ValueError(
-                "Datasource ID not found. Please include it as a dbt variable in the dbt_project.yml. Instructions: {cloud_doc_url}"
-                "\nvars:\n data_diff:\n   datasource_id: 1234"
+            raise DataDiffNoDatasourceIdError(
+                f"Datasource ID not found. Please include it as a dbt variable in the dbt_project.yml. \nInstructions: {CLOUD_DOC_URL}\n\nvars:\n data_diff:\n   datasource_id: 1234"
             )
 
         data_source = api.get_data_source(config.datasource_id)
@@ -283,8 +281,9 @@ def _initialize_api() -> Optional[DatafoldAPI]:
         rich.print("[red]API key not found. Getting from the keyring service")
         api_key = keyring.get_password("data-diff", "DATAFOLD_API_KEY")
         if not api_key:
-            cloud_doc_url = "https://docs.datafold.com/development_testing/cloud"
-            rich.print("[red]API key not found. Please follow the steps at {cloud_doc_url} to use the --cloud flag.")
+            raise DataDiffNoAPIKeyError(
+                f"API key not found. Please follow the steps at {CLOUD_DOC_URL} to use the --cloud flag."
+            )
     rich.print("Saving the API key to the system keyring service")
     try:
         keyring.set_password("data-diff", "DATAFOLD_API_KEY", api_key)

data_diff/errors.py

@@ -48,3 +48,9 @@ class DataDiffDbtSelectVersionTooLowError(Exception):
 
 class DataDiffCustomSchemaNoConfigError(Exception):
     "Raised when a model has a custom schema, but there is no prod_custom_schema config. (And not using --state)."
+
+class DataDiffNoAPIKeyError(Exception):
+    "Raised when using --cloud but no API key is present in the DATAFOLD_API_KEY env var or keyring"
+
+class DataDiffNoDatasourceIdError(Exception):
+    "Raised when using --cloud but no datasource_id was found in dbt_project.yml"

tests/test_dbt.py

@@ -16,6 +16,8 @@
     DataDiffDbtRunResultsVersionError,
     DataDiffDbtSelectVersionTooLowError,
     DataDiffDbtSnowflakeSetConnectionError,
+    DataDiffNoAPIKeyError,
+    DataDiffNoDatasourceIdError,
 )
 from .test_cli import run_datadiff_cli
 
@@ -771,7 +773,7 @@ def test_diff_is_cloud_no_ds_id(
         )
         mock_get_diff_vars.return_value = diff_vars
 
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DataDiffNoDatasourceIdError):
             dbt_diff(is_cloud=True)
         mock_dbt_parser_inst.get_models.assert_called_once()
         mock_dbt_parser_inst.set_connection.assert_not_called()
@@ -781,6 +783,31 @@ def test_diff_is_cloud_no_ds_id(
         mock_local_diff.assert_not_called()
         mock_print.assert_called_once()
 
+    @patch("data_diff.dbt._get_diff_vars")
+    @patch("data_diff.dbt._local_diff")
+    @patch("data_diff.dbt._cloud_diff")
+    @patch("data_diff.dbt_parser.DbtParser.__new__")
+    @patch("builtins.input", return_value="n")
+    def test_diff_is_cloud_no_api_key(
+        self, _, mock_dbt_parser, mock_cloud_diff, mock_local_diff, mock_get_diff_vars
+    ):
+        mock_dbt_parser_inst = Mock()
+        mock_model = Mock()
+        config = TDatadiffConfig(prod_database="prod_db", prod_schema="prod_schema")
+
+        mock_dbt_parser.return_value = mock_dbt_parser_inst
+        mock_dbt_parser_inst.get_models.return_value = [mock_model]
+        mock_dbt_parser_inst.get_datadiff_config.return_value = config
+
+        with self.assertRaises(DataDiffNoAPIKeyError):
+            dbt_diff(is_cloud=True)
+        mock_dbt_parser_inst.get_models.assert_called_once()
+        mock_dbt_parser_inst.set_connection.assert_not_called()
+
+        mock_get_diff_vars.assert_not_called()
+        mock_cloud_diff.assert_not_called()
+        mock_local_diff.assert_not_called()
+
     @patch("data_diff.dbt._get_diff_vars")
     @patch("data_diff.dbt._local_diff")
     @patch("data_diff.dbt._cloud_diff")