fix: resolve issue in Cache.create_source_tables (airbytehq#642)

Author: aaronsteers

airbyte/caches/base.py

@@ -313,6 +313,13 @@ def create_source_tables(
 
         catalog_provider = CatalogProvider(source.get_configured_catalog(streams=streams))
 
+        # Register the incoming source catalog
+        self.register_source(
+            source_name=source.name,
+            incoming_source_catalog=catalog_provider.configured_catalog,
+            stream_names=set(catalog_provider.stream_names),
+        )
+
         # Ensure schema exists
         self.processor._ensure_schema_exists()  # noqa: SLF001  # Accessing non-public member
 

tests/integration_tests/test_all_cache_types.py

@@ -200,6 +200,19 @@ def test_replace_strategy(
     assert len(list(result.cache.streams["users"])) == FAKER_SCALE_A
 
 
+@pytest.mark.requires_creds
+@pytest.mark.slow
+def test_cache_create_source_tables(
+    source_faker_seed_a: ab.Source,
+    new_generic_cache: ab.caches.CacheBase,
+) -> None:
+    """Test that the cache creation and source tables work as expected."""
+    new_generic_cache.create_source_tables(source_faker_seed_a)
+    assert set(new_generic_cache.streams.keys()) == set(
+        source_faker_seed_a.get_selected_streams()
+    )
+
+
 @pytest.mark.requires_creds
 @pytest.mark.slow
 def test_merge_strategy(

tests/unit_tests/test_caches.py

@@ -61,77 +61,3 @@ def test_duck_db_cache_config_get_database_name_with_default_schema_name():
 
 def test_duck_db_cache_config_inheritance_from_sql_cache_config_base():
     assert issubclass(DuckDBCache, CacheBase)
-
-
-def test_create_source_tables(mocker):
-    """Test that the create_source_tables method correctly creates tables based on the source's catalog."""
-    # Import here to avoid circular imports
-    from airbyte_protocol.models import (
-        ConfiguredAirbyteCatalog,
-        ConfiguredAirbyteStream,
-    )
-
-    # Create a proper ConfiguredAirbyteCatalog for mocking
-    stream1 = ConfiguredAirbyteStream(
-        stream={
-            "name": "stream1",
-            "json_schema": {},
-            "supported_sync_modes": ["full_refresh"],
-        },
-        sync_mode="full_refresh",
-        destination_sync_mode="overwrite",
-    )
-    stream2 = ConfiguredAirbyteStream(
-        stream={
-            "name": "stream2",
-            "json_schema": {},
-            "supported_sync_modes": ["full_refresh"],
-        },
-        sync_mode="full_refresh",
-        destination_sync_mode="overwrite",
-    )
-    catalog = ConfiguredAirbyteCatalog(streams=[stream1, stream2])
-
-    # Mock the catalog provider
-    mock_catalog_provider = mocker.Mock()
-    mock_catalog_provider.stream_names = ["stream1", "stream2"]
-    mocker.patch(
-        "airbyte.shared.catalog_providers.CatalogProvider",
-        return_value=mock_catalog_provider,
-    )
-
-    # Mock a source with configured catalog and selected streams
-    mock_source = mocker.Mock()
-    mock_source.get_configured_catalog.return_value = catalog
-    mock_source.get_selected_streams.return_value = ["stream1"]
-
-    # Create a DuckDBCache instance with mocked processor
-    cache = DuckDBCache(db_path=UNIT_TEST_DB_PATH)
-
-    # Mock the processor property
-    mock_processor = mocker.Mock()
-    mocker.patch.object(
-        DuckDBCache, "processor", mocker.PropertyMock(return_value=mock_processor)
-    )
-
-    # Test with default (None) stream parameter - should use source's selected streams
-    cache.create_source_tables(mock_source)
-
-    # Verify the correct methods were called
-    mock_source.get_selected_streams.assert_called_once()
-    mock_source.get_configured_catalog.assert_called_once_with(streams=["stream1"])
-    mock_processor._ensure_schema_exists.assert_called_once()
-    assert mock_processor._ensure_final_table_exists.call_count == 2
-
-    # Reset mocks
-    mock_source.reset_mock()
-    mock_processor.reset_mock()
-
-    # Test with explicit stream list
-    cache.create_source_tables(mock_source, streams=["stream2"])
-
-    # Verify the correct methods were called
-    mock_source.get_selected_streams.assert_not_called()
-    mock_source.get_configured_catalog.assert_called_once_with(streams=["stream2"])
-    mock_processor._ensure_schema_exists.assert_called_once()
-    assert mock_processor._ensure_final_table_exists.call_count == 2