file-mode-api: Mock the test instead of messing with Classes

Author: aldogonzalez8

airbyte_cdk/sources/declarative/yaml_declarative_source.py

@@ -24,7 +24,6 @@ def __init__(
         catalog: Optional[ConfiguredAirbyteCatalog] = None,
         config: Optional[Mapping[str, Any]] = None,
         state: Optional[List[AirbyteStateMessage]] = None,
-        emit_connector_builder_messages: Optional[bool] = False,
     ) -> None:
         """
         :param path_to_yaml: Path to the yaml file describing the source
@@ -36,8 +35,7 @@ def __init__(
             catalog=catalog or ConfiguredAirbyteCatalog(streams=[]),
             config=config or {},
             state=state or [],
-            source_config=source_config,
-            emit_connector_builder_messages=emit_connector_builder_messages or False,
+            source_config=source_config
         )
 
     def _read_and_parse_yaml_file(self, path_to_yaml_file: str) -> ConnectionDefinition:

unit_tests/sources/declarative/file/test_file_stream.py

@@ -3,9 +3,10 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 from unittest import TestCase
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 from airbyte_cdk.models import AirbyteStateMessage, ConfiguredAirbyteCatalog, Status
+from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ModelToComponentFactory as OriginalModelToComponentFactory
 from airbyte_cdk.sources.declarative.retrievers.file_uploader.noop_file_writer import NoopFileWriter
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
 from airbyte_cdk.test.catalog_builder import CatalogBuilder, ConfiguredAirbyteStreamBuilder
@@ -35,16 +36,14 @@ def _source(
     config: Dict[str, Any],
     state: Optional[List[AirbyteStateMessage]] = None,
     yaml_file: Optional[str] = None,
-    emit_connector_builder_messages: Optional[bool] = False,
 ) -> YamlDeclarativeSource:
     if not yaml_file:
         yaml_file = "file_stream_manifest.yaml"
     return YamlDeclarativeSource(
         path_to_yaml=str(Path(__file__).parent / yaml_file),
         catalog=catalog,
         config=config,
-        state=state,
-        emit_connector_builder_messages=emit_connector_builder_messages,
+        state=state
     )
 
 
@@ -53,13 +52,12 @@ def read(
     catalog: ConfiguredAirbyteCatalog,
     state_builder: Optional[StateBuilder] = None,
     expecting_exception: bool = False,
-    yaml_file: Optional[str] = None,
-    emit_connector_builder_messages: Optional[bool] = False,
+    yaml_file: Optional[str] = None
 ) -> EntrypointOutput:
     config = config_builder.build()
     state = state_builder.build() if state_builder else StateBuilder().build()
     return entrypoint_read(
-        _source(catalog, config, state, yaml_file, emit_connector_builder_messages),
+        _source(catalog, config, state, yaml_file),
         config,
         catalog,
         state,
@@ -185,7 +183,7 @@ def test_get_article_attachments_with_filename_extractor(self) -> None:
                 yaml_file="test_file_stream_with_filename_extractor.yaml",
             )
 
-            assert output.records
+            assert len(output.records) == 1
             file_reference = output.records[0].record.file_reference
             assert file_reference
             assert (
@@ -217,30 +215,40 @@ def test_get_article_attachments_messages_for_connector_builder(self) -> None:
                 ),
             )
 
-            output = read(
-                self._config(),
-                CatalogBuilder()
-                .with_stream(ConfiguredAirbyteStreamBuilder().with_name("article_attachments"))
-                .build(),
-                yaml_file="test_file_stream_with_filename_extractor.yaml",
-                emit_connector_builder_messages=True,
-            )
+            # Define a mock factory that forces emit_connector_builder_messages=True
+            class MockModelToComponentFactory(OriginalModelToComponentFactory):
+                def __init__(self, *args, **kwargs):
+                    kwargs['emit_connector_builder_messages'] = True
+                    super().__init__(*args, **kwargs)
 
-            assert len(output.records) == 1
-            file_reference = output.records[0].record.file_reference
-            assert file_reference
-            assert file_reference.staging_file_url
-            assert file_reference.source_file_relative_path
-            # because we didn't write the file, the size is 0
-            assert file_reference.file_size_bytes == NoopFileWriter.NOOP_FILE_SIZE
+            # Patch the factory class where ConcurrentDeclarativeSource (parent of YamlDeclarativeSource) imports it
+            with patch(
+                "airbyte_cdk.sources.declarative.concurrent_declarative_source.ModelToComponentFactory",
+                new=MockModelToComponentFactory
+            ):
+                output = read(
+                    self._config(),
+                    CatalogBuilder()
+                    .with_stream(ConfiguredAirbyteStreamBuilder().with_name("article_attachments"))
+                    .build(),
+                    yaml_file="test_file_stream_with_filename_extractor.yaml"
+                )
 
-            # Assert file reference fields are copied to record data
-            record_data = output.records[0].record.data
-            assert record_data["staging_file_url"] == file_reference.staging_file_url
-            assert (
-                record_data["source_file_relative_path"] == file_reference.source_file_relative_path
-            )
-            assert record_data["file_size_bytes"] == file_reference.file_size_bytes
+                assert len(output.records) == 1
+                file_reference = output.records[0].record.file_reference
+                assert file_reference
+                assert file_reference.staging_file_url
+                assert file_reference.source_file_relative_path
+                # because we didn't write the file, the size is NOOP_FILE_SIZE
+                assert file_reference.file_size_bytes == NoopFileWriter.NOOP_FILE_SIZE
+
+                # Assert file reference fields are copied to record data
+                record_data = output.records[0].record.data
+                assert record_data["staging_file_url"] == file_reference.staging_file_url
+                assert (
+                    record_data["source_file_relative_path"] == file_reference.source_file_relative_path
+                )
+                assert record_data["file_size_bytes"] == file_reference.file_size_bytes
 
     def test_discover_article_attachments(self) -> None:
         output = discover(self._config())