66from unittest .mock import Mock , patch
77
88from airbyte_cdk .models import AirbyteStateMessage , ConfiguredAirbyteCatalog , Status
9- from airbyte_cdk .sources .declarative .parsers .model_to_component_factory import ModelToComponentFactory as OriginalModelToComponentFactory
9+ from airbyte_cdk .sources .declarative .parsers .model_to_component_factory import (
10+ ModelToComponentFactory as OriginalModelToComponentFactory ,
11+ )
1012from airbyte_cdk .sources .declarative .retrievers .file_uploader .noop_file_writer import NoopFileWriter
1113from airbyte_cdk .sources .declarative .yaml_declarative_source import YamlDeclarativeSource
1214from airbyte_cdk .test .catalog_builder import CatalogBuilder , ConfiguredAirbyteStreamBuilder
@@ -43,7 +45,7 @@ def _source(
4345 path_to_yaml = str (Path (__file__ ).parent / yaml_file ),
4446 catalog = catalog ,
4547 config = config ,
46- state = state
48+ state = state ,
4749 )
4850
4951
@@ -52,7 +54,7 @@ def read(
5254 catalog : ConfiguredAirbyteCatalog ,
5355 state_builder : Optional [StateBuilder ] = None ,
5456 expecting_exception : bool = False ,
55- yaml_file : Optional [str ] = None
57+ yaml_file : Optional [str ] = None ,
5658) -> EntrypointOutput :
5759 config = config_builder .build ()
5860 state = state_builder .build () if state_builder else StateBuilder ().build ()
@@ -218,20 +220,20 @@ def test_get_article_attachments_messages_for_connector_builder(self) -> None:
218220 # Define a mock factory that forces emit_connector_builder_messages=True
219221 class MockModelToComponentFactory (OriginalModelToComponentFactory ):
220222 def __init__ (self , * args , ** kwargs ):
221- kwargs [' emit_connector_builder_messages' ] = True
223+ kwargs [" emit_connector_builder_messages" ] = True
222224 super ().__init__ (* args , ** kwargs )
223225
224226 # Patch the factory class where ConcurrentDeclarativeSource (parent of YamlDeclarativeSource) imports it
225227 with patch (
226228 "airbyte_cdk.sources.declarative.concurrent_declarative_source.ModelToComponentFactory" ,
227- new = MockModelToComponentFactory
229+ new = MockModelToComponentFactory ,
228230 ):
229231 output = read (
230232 self ._config (),
231233 CatalogBuilder ()
232234 .with_stream (ConfiguredAirbyteStreamBuilder ().with_name ("article_attachments" ))
233235 .build (),
234- yaml_file = "test_file_stream_with_filename_extractor.yaml"
236+ yaml_file = "test_file_stream_with_filename_extractor.yaml" ,
235237 )
236238
237239 assert len (output .records ) == 1
@@ -246,7 +248,8 @@ def __init__(self, *args, **kwargs):
246248 record_data = output .records [0 ].record .data
247249 assert record_data ["staging_file_url" ] == file_reference .staging_file_url
248250 assert (
249- record_data ["source_file_relative_path" ] == file_reference .source_file_relative_path
251+ record_data ["source_file_relative_path" ]
252+ == file_reference .source_file_relative_path
250253 )
251254 assert record_data ["file_size_bytes" ] == file_reference .file_size_bytes
252255
0 commit comments