run black

dongwonmoon · dongwonmoon · commit f8ea7147ab9b · 2025-11-11T16:03:40.000+09:00
diff --git a/schema_scribe/components/writers/notion_writer.py b/schema_scribe/components/writers/notion_writer.py
@@ -53,20 +53,28 @@ def _connect(self):
         api_token_param = self.params.get("api_token")
         token_to_use = None
 
-        logger.debug(f"NotionWriter._connect: raw api_token_param: {api_token_param}")
+        logger.debug(
+            f"NotionWriter._connect: raw api_token_param: {api_token_param}"
+        )
 
         if api_token_param:
-            if api_token_param.startswith("${") and api_token_param.endswith("}"):
+            if api_token_param.startswith("${") and api_token_param.endswith(
+                "}"
+            ):
                 env_var = api_token_param[2:-1]
                 token_to_use = os.getenv(env_var)
-                logger.debug(f"NotionWriter._connect: resolved from env '{env_var}': {token_to_use}")
+                logger.debug(
+                    f"NotionWriter._connect: resolved from env '{env_var}': {token_to_use}"
+                )
                 if not token_to_use:
                     raise ConfigError(
                         f"Environment variable '{env_var}' is required but not set."
                     )
             else:
                 token_to_use = api_token_param
-                logger.debug(f"NotionWriter._connect: using direct token: {token_to_use}")
+                logger.debug(
+                    f"NotionWriter._connect: using direct token: {token_to_use}"
+                )
 
         if not token_to_use:
             raise ConfigError(
@@ -137,23 +145,23 @@ def _H2(self, text: str) -> Dict[str, Any]:
         return {
             "object": "block",
             "type": "heading_2",
-            "heading_2": {"rich_text": [{"text": {"content": text}}]}
+            "heading_2": {"rich_text": [{"text": {"content": text}}]},
         }
 
     def _H3(self, text: str) -> Dict[str, Any]:
         """Creates a Notion Heading 3 block."""
         return {
             "object": "block",
             "type": "heading_3",
-            "heading_3": {"rich_text": [{"text": {"content": text}}]}
+            "heading_3": {"rich_text": [{"text": {"content": text}}]},
         }
 
     def _Para(self, text: str) -> Dict[str, Any]:
         """Creates a Notion Paragraph block."""
         return {
             "object": "block",
             "type": "paragraph",
-            "paragraph": {"rich_text": [{"text": {"content": text}}]}
+            "paragraph": {"rich_text": [{"text": {"content": text}}]},
         }
 
     def _Code(self, text: str, lang: str = "sql") -> Dict[str, Any]:
@@ -176,8 +184,7 @@ def _clean_mermaid_code(self, code: str) -> str:
         return code.replace("```mermaid", "").replace("```", "").strip()
 
     def _generate_notion_blocks(
-        self,
-        catalog_data: Dict[str, Any]
+        self, catalog_data: Dict[str, Any]
     ) -> List[Dict[str, Any]]:
         """
         Dynamically generates Notion blocks by detecting the catalog structure.
@@ -215,9 +222,7 @@ def _generate_notion_blocks(
             return [self._Para("Unknown catalog structure provided.")]
 
     def _create_column_table(
-        self,
-        columns: List[Dict[str, Any]],
-        is_dbt: bool = False
+        self, columns: List[Dict[str, Any]], is_dbt: bool = False
     ) -> Dict[str, Any]:
         """
         Creates a Notion Table block to display column details.
@@ -287,8 +292,7 @@ def _generate_mermaid_erd(self, foreign_keys: List[Dict[str, str]]) -> str:
         return "\n".join(code)
 
     def _generate_db_blocks(
-        self,
-        catalog_data: Dict[str, Any]
+        self, catalog_data: Dict[str, Any]
     ) -> List[Dict[str, Any]]:
         """
         Generates a list of Notion blocks for a traditional database catalog.
@@ -334,8 +338,7 @@ def _generate_db_blocks(
         return blocks
 
     def _generate_dbt_blocks(
-        self,
-        catalog_data: Dict[str, Any]
+        self, catalog_data: Dict[str, Any]
     ) -> List[Dict[str, Any]]:
         """
         Generates a list of Notion blocks for a dbt project catalog.
diff --git a/schema_scribe/services/catalog_generator.py b/schema_scribe/services/catalog_generator.py
@@ -217,4 +217,4 @@ def generate_catalog(
         catalog_data["foreign_keys"] = foreign_keys
 
         logger.info("Catalog generation completed.")
-        return catalog_data
+        return catalog_data
diff --git a/schema_scribe/tests/integration/test_db_workflow.py b/schema_scribe/tests/integration/test_db_workflow.py
@@ -11,7 +11,11 @@
 from unittest.mock import MagicMock, patch
 
 from schema_scribe.workflows.db_workflow import DbWorkflow
-from schema_scribe.core.interfaces import BaseConnector, BaseLLMClient, BaseWriter
+from schema_scribe.core.interfaces import (
+    BaseConnector,
+    BaseLLMClient,
+    BaseWriter,
+)
 
 
 @pytest.fixture
@@ -27,33 +31,97 @@ def mock_db_connector(sqlite_db: str):
         {"name": "orders", "comment": None},
     ]
     mock_connector.get_views.return_value = [
-        {"name": "user_orders", "definition": "SELECT * FROM users JOIN orders ON users.id = orders.user_id"}
+        {
+            "name": "user_orders",
+            "definition": "SELECT * FROM users JOIN orders ON users.id = orders.user_id",
+        }
     ]
     mock_connector.get_columns.side_effect = [
         # users table columns
         [
-            {"name": "id", "type": "INTEGER", "is_nullable": False, "is_pk": True, "comment": None},
-            {"name": "email", "type": "TEXT", "is_nullable": False, "is_pk": False, "comment": None},
+            {
+                "name": "id",
+                "type": "INTEGER",
+                "is_nullable": False,
+                "is_pk": True,
+                "comment": None,
+            },
+            {
+                "name": "email",
+                "type": "TEXT",
+                "is_nullable": False,
+                "is_pk": False,
+                "comment": None,
+            },
         ],
         # products table columns
         [
-            {"name": "id", "type": "INTEGER", "is_nullable": False, "is_pk": True, "comment": None},
-            {"name": "name", "type": "TEXT", "is_nullable": False, "is_pk": False, "comment": None},
-            {"name": "price", "type": "REAL", "is_nullable": False, "is_pk": False, "comment": None},
+            {
+                "name": "id",
+                "type": "INTEGER",
+                "is_nullable": False,
+                "is_pk": True,
+                "comment": None,
+            },
+            {
+                "name": "name",
+                "type": "TEXT",
+                "is_nullable": False,
+                "is_pk": False,
+                "comment": None,
+            },
+            {
+                "name": "price",
+                "type": "REAL",
+                "is_nullable": False,
+                "is_pk": False,
+                "comment": None,
+            },
         ],
         # orders table columns
         [
-            {"name": "id", "type": "INTEGER", "is_nullable": False, "is_pk": True, "comment": None},
-            {"name": "user_id", "type": "INTEGER", "is_nullable": False, "is_pk": False, "comment": None},
-            {"name": "product_id", "type": "INTEGER", "is_nullable": False, "is_pk": False, "comment": None},
+            {
+                "name": "id",
+                "type": "INTEGER",
+                "is_nullable": False,
+                "is_pk": True,
+                "comment": None,
+            },
+            {
+                "name": "user_id",
+                "type": "INTEGER",
+                "is_nullable": False,
+                "is_pk": False,
+                "comment": None,
+            },
+            {
+                "name": "product_id",
+                "type": "INTEGER",
+                "is_nullable": False,
+                "is_pk": False,
+                "comment": None,
+            },
         ],
     ]
     mock_connector.get_foreign_keys.return_value = [
-        {"source_table": "orders", "source_column": "user_id", "target_table": "users", "target_column": "id"},
-        {"source_table": "orders", "source_column": "product_id", "target_table": "products", "target_column": "id"},
+        {
+            "source_table": "orders",
+            "source_column": "user_id",
+            "target_table": "users",
+            "target_column": "id",
+        },
+        {
+            "source_table": "orders",
+            "source_column": "product_id",
+            "target_table": "products",
+            "target_column": "id",
+        },
     ]
     mock_connector.get_column_profile.return_value = {
-        "total_count": 0, "null_count": 0, "distinct_count": 0, "is_unique": True
+        "total_count": 0,
+        "null_count": 0,
+        "distinct_count": 0,
+        "is_unique": True,
     }
     mock_connector.close.return_value = None
     return mock_connector
@@ -66,7 +134,9 @@ def mock_llm_client():
     """
     mock_client = MagicMock(spec=BaseLLMClient)
     mock_client.llm_profile_name = "test_llm"
-    mock_client.get_description.return_value = "This is an AI-generated description."
+    mock_client.get_description.return_value = (
+        "This is an AI-generated description."
+    )
     return mock_client
 
 
@@ -113,9 +183,11 @@ def test_db_workflow_end_to_end(
     # 1. Verify components were called
     mock_db_connector.get_tables.assert_called_once()
     mock_db_connector.get_views.assert_called_once()
-    assert mock_db_connector.get_columns.call_count == 3 # Only called for tables
+    assert (
+        mock_db_connector.get_columns.call_count == 3
+    )  # Only called for tables
     mock_db_connector.get_foreign_keys.assert_called_once()
-    mock_llm_client.get_description.assert_called() # Called for tables, views, columns
+    mock_llm_client.get_description.assert_called()  # Called for tables, views, columns
     mock_writer.write.assert_called_once()
     mock_db_connector.close.assert_called_once()
 
@@ -129,7 +201,7 @@ def test_db_workflow_end_to_end(
 
     assert captured_writer_params == {
         "db_profile_name": "test_db",
-        "db_connector": mock_db_connector, # Added mock_db_connector to expected params
+        "db_connector": mock_db_connector,  # Added mock_db_connector to expected params
         "output_filename": str(output_md_path),
     }
 
@@ -182,7 +254,11 @@ def test_db_workflow_end_to_end_with_profiling(
     for call in calls:
         prompt_text = call[0][0]
         # Make the search more flexible
-        if "Table: {'name': 'users', 'comment': None}" in prompt_text and "Column: id" in prompt_text and "Data Profile Context:" in prompt_text:
+        if (
+            "Table: {'name': 'users', 'comment': None}" in prompt_text
+            and "Column: id" in prompt_text
+            and "Data Profile Context:" in prompt_text
+        ):
             users_id_prompt = prompt_text
             break
 
diff --git a/schema_scribe/tests/integration/test_dbt_workflow.py b/schema_scribe/tests/integration/test_dbt_workflow.py
@@ -16,7 +16,11 @@
 
 from schema_scribe.workflows.dbt_workflow import DbtWorkflow
 from schema_scribe.core.exceptions import CIError
-from schema_scribe.core.interfaces import BaseConnector, BaseLLMClient, BaseWriter
+from schema_scribe.core.interfaces import (
+    BaseConnector,
+    BaseLLMClient,
+    BaseWriter,
+)
 
 # A minimal manifest.json structure needed for the tests
 MINIMAL_MANIFEST = {
@@ -129,7 +133,9 @@ def mock_llm_client():
     """
     mock_client = MagicMock(spec=BaseLLMClient)
     mock_client.llm_profile_name = "test_llm"
-    mock_client.get_description.return_value = "This is an AI-generated description."
+    mock_client.get_description.return_value = (
+        "This is an AI-generated description."
+    )
     return mock_client
 
 
@@ -221,7 +227,9 @@ def test_dbt_workflow_update(
                 {"name": "first_name", "description": "", "type": "TEXT"},
             ],
             "path": "customers.sql",
-            "original_file_path": os.path.join(dbt_project, "models", "customers.sql"),
+            "original_file_path": os.path.join(
+                dbt_project, "models", "customers.sql"
+            ),
         }
     ]
     mock_parser_constructor.return_value = mock_parser_instance
@@ -282,7 +290,9 @@ def test_dbt_workflow_check_fails(
                 {"name": "first_name", "description": "", "type": "TEXT"},
             ],
             "path": "customers.sql",
-            "original_file_path": os.path.join(dbt_project, "models", "customers.sql"),
+            "original_file_path": os.path.join(
+                dbt_project, "models", "customers.sql"
+            ),
         }
     ]
     mock_parser_constructor.return_value = mock_parser_instance
@@ -327,7 +337,9 @@ def test_dbt_workflow_check_succeeds(
                 {"name": "first_name", "description": "", "type": "TEXT"},
             ],
             "path": "customers.sql",
-            "original_file_path": os.path.join(dbt_project, "models", "customers.sql"),
+            "original_file_path": os.path.join(
+                dbt_project, "models", "customers.sql"
+            ),
         }
     ]
     mock_parser_constructor.return_value = mock_parser_instance
@@ -501,4 +513,4 @@ def test_dbt_workflow_drift_mode_drift_detected(
     assert (
         "Unique ID for the customer" in drift_prompt_call[0][0]
     )  # Existing doc
-    assert "Is Unique: False" in drift_prompt_call[0][0]  # Conflicting stat
+    assert "Is Unique: False" in drift_prompt_call[0][0]  # Conflicting stat
diff --git a/schema_scribe/tests/integration/test_lineage_workflow.py b/schema_scribe/tests/integration/test_lineage_workflow.py
@@ -36,7 +36,7 @@ def test_lineage_workflow_e2e(
             "target_column": "id",
         }
     ]
-    mock_db_connector.db_profile_name = "test_db" # Add this for logging
+    mock_db_connector.db_profile_name = "test_db"  # Add this for logging
     mock_db_connector.close.return_value = None
 
     # --- Mock dbt Parser (Logical Lineage) ---
diff --git a/schema_scribe/tests/unit/llm_clients/test_google_client.py b/schema_scribe/tests/unit/llm_clients/test_google_client.py
@@ -51,5 +51,5 @@ def test_google_client_get_description(mock_genai, mocker):
     # The actual call includes generation_config
     mock_model_instance.generate_content.assert_called_once_with(
         "test prompt",
-        generation_config=mock_genai.GenerationConfig(max_output_tokens=150)
-    )
+        generation_config=mock_genai.GenerationConfig(max_output_tokens=150),
+    )
diff --git a/schema_scribe/tests/unit/writers/test_notion_writer.py b/schema_scribe/tests/unit/writers/test_notion_writer.py
@@ -121,6 +121,7 @@ def test_notion_writer_success(mock_notion_client, mock_db_catalog_data):
     )
     assert table_h2 is not None
 
+
 @patch.dict(os.environ, {"NOTION_TEST_KEY": "env_key_value"})
 @patch("schema_scribe.components.writers.notion_writer.Client")
 def test_notion_writer_resolves_env_var(
@@ -134,7 +135,7 @@ def test_notion_writer_resolves_env_var(
     the Notion client.
     """
     mock_notion_client.return_value = MagicMock()  # Basic mock
-    
+
     writer = NotionWriter()
     kwargs = {
         "api_token": "${NOTION_TEST_KEY}",  # Reference the env var
@@ -143,7 +144,6 @@ def test_notion_writer_resolves_env_var(
 
     writer.write(mock_db_catalog_data, **kwargs)
 
-    
     # Assert client was initialized with the *resolved* key
     mock_notion_client.assert_called_once()
     call_args, call_kwargs = mock_notion_client.call_args
@@ -165,4 +165,4 @@ def test_notion_writer_config_errors(mock_db_catalog_data):
 
     # 2. Missing parent_page_id
     with pytest.raises(ConfigError, match="'parent_page_id' is required"):
-        writer.write(mock_db_catalog_data, api_token="fake_token")
+        writer.write(mock_db_catalog_data, api_token="fake_token")

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ def test_lineage_workflow_e2e(`
`36`	`36`	`"target_column": "id",`
`37`	`37`	`}`
`38`	`38`	`]`
`39`		`- mock_db_connector.db_profile_name = "test_db" # Add this for logging`
	`39`	`+ mock_db_connector.db_profile_name = "test_db" # Add this for logging`
`40`	`40`	`mock_db_connector.close.return_value = None`
`41`	`41`
`42`	`42`	`# --- Mock dbt Parser (Logical Lineage) ---`