Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -491,7 +491,7 @@ def _uses_default_streams(self) -> bool:
return type(self).streams is ConcurrentDeclarativeSource.streams

def _streams_by_name(self, stream_names: List[str]) -> Optional[List[AbstractStream]]:
"""Build the named static streams plus every dynamic stream.
"""Build the named static streams without resolving or building dynamic streams.

Return None when the names are empty, any name isn't static, or `streams()` is overridden.
"""
Expand All @@ -506,14 +506,13 @@ def _streams_by_name(self, stream_names: List[str]) -> Optional[List[AbstractStr
return None

self._prepare_for_stream_build()
dynamic_configs = self.dynamic_streams
selected_names = set(stream_names)
selected_configs = [
stream_config
for stream_config in static_configs
if stream_config.get("name") in selected_names
] + dynamic_configs
return self._create_streams(selected_configs, static_configs + dynamic_configs)
]
return self._create_streams(selected_configs, static_configs)

def _first_streams(self, count: int) -> Optional[List[AbstractStream]]:
"""Build the first `count` streams, static then dynamic, in `streams()` order.
Expand Down
133 changes: 128 additions & 5 deletions unit_tests/sources/declarative/checks/test_check_stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from airbyte_cdk.sources.streams.core import Stream
from airbyte_cdk.sources.streams.http import HttpStream
from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
from airbyte_cdk.utils import AirbyteTracedException

logger = logging.getLogger("test")
config = dict()
Expand Down Expand Up @@ -1113,6 +1114,7 @@ def test_check_for_dynamic_stream_name_falls_back_to_streams(mocker):
state=None,
)
streams_spy = mocker.spy(source, "streams")
dynamic_configs_spy = mocker.spy(source, "_dynamic_stream_configs")

with HttpMocker() as http_mocker:
http_mocker.get(
Expand All @@ -1126,9 +1128,92 @@ def test_check_for_dynamic_stream_name_falls_back_to_streams(mocker):
assert source.check(logger, source._config).status == Status.SUCCEEDED

streams_spy.assert_called_once()
dynamic_configs_spy.assert_called_once()


def test_check_static_stream_does_not_resolve_http_dynamic_stream():
manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT)
manifest["dynamic_streams"] = [manifest["dynamic_streams"][0]]
manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]}
source = ConcurrentDeclarativeSource(
source_config=manifest,
config={**_CONFIG, "api_key": "test"},
catalog=None,
state=None,
)

with HttpMocker() as http_mocker:
http_mocker.get(
HttpRequest(url="https://api.test.com/static"),
HttpResponse(body=json.dumps([])),
)
assert source.check(logger, source._config).status == Status.SUCCEEDED


def test_check_with_duplicate_dynamic_streams_skips_resolution():
manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT)
manifest["dynamic_streams"] = [manifest["dynamic_streams"][1]]
manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]}
config_with_duplicate_dynamic_streams = {
**_CONFIG,
"api_key": "test",
"custom_streams": [
{"id": 1, "name": "duplicate"},
{"id": 2, "name": "duplicate"},
],
}
source = ConcurrentDeclarativeSource(
source_config=manifest,
config=config_with_duplicate_dynamic_streams,
catalog=None,
state=None,
)

with HttpMocker() as http_mocker:
http_mocker.get(
HttpRequest(url="https://api.test.com/static"),
HttpResponse(body=json.dumps([])),
)
assert source.check(logger, source._config).status == Status.SUCCEEDED

with pytest.raises(
AirbyteTracedException,
match="Dynamic streams list contains a duplicate name: duplicate. Please check your configuration.",
):
source.discover(logger=source.logger, config=source._config)


def test_check_skips_http_dynamic_stream_resolver_errors():
manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT)
manifest["dynamic_streams"] = [manifest["dynamic_streams"][0]]
manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]}
source = ConcurrentDeclarativeSource(
source_config=manifest,
config={**_CONFIG, "api_key": "test"},
catalog=None,
state=None,
)
resolver_request = HttpRequest(url="https://api.test.com/items")

with HttpMocker() as http_mocker:
http_mocker.get(
HttpRequest(url="https://api.test.com/static"),
HttpResponse(body=json.dumps([])),
)
http_mocker.get(
resolver_request,
HttpResponse(body=json.dumps({"error": "server error"}), status_code=500),
)
assert source.check(logger, source._config).status == Status.SUCCEEDED
http_mocker.assert_number_of_calls(resolver_request, 0)

with pytest.raises(AirbyteTracedException, match="HTTP Status Code: 500"):
source.discover(logger=source.logger, config=source._config)

def test_check_dynamic_stream_configs_are_evaluated_once_when_not_checked(mocker):

def test_check_dynamic_stream_configs_are_not_evaluated_when_only_static_streams_are_checked(
mocker,
):
manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT)
manifest["dynamic_streams"] = [manifest["dynamic_streams"][1]]
manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]}
Expand All @@ -1147,7 +1232,43 @@ def test_check_dynamic_stream_configs_are_evaluated_once_when_not_checked(mocker
dynamic_configs_spy = mocker.spy(source, "_dynamic_stream_configs")
assert source.check(logger, source._config).status == Status.SUCCEEDED

dynamic_configs_spy.assert_called_once()
dynamic_configs_spy.assert_not_called()


def test_check_with_streams_override_resolves_http_dynamic_streams():
class StreamsOverrideSource(ConcurrentDeclarativeSource):
def __init__(self, *args, **kwargs):
self.streams_call_count = 0
super().__init__(*args, **kwargs)

def streams(self, config):
self.streams_call_count += 1
return super().streams(config=config)

manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT)
manifest["dynamic_streams"] = [manifest["dynamic_streams"][0]]
manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]}
source = StreamsOverrideSource(
source_config=manifest,
config={**_CONFIG, "api_key": "test"},
catalog=None,
state=None,
)
resolver_request = HttpRequest(url="https://api.test.com/items")

with HttpMocker() as http_mocker:
http_mocker.get(
HttpRequest(url="https://api.test.com/static"),
HttpResponse(body=json.dumps([])),
)
http_mocker.get(
resolver_request,
HttpResponse(body=json.dumps([{"id": 3, "name": "item_3"}])),
)
assert source.check(logger, source._config).status == Status.SUCCEEDED
http_mocker.assert_number_of_calls(resolver_request, 1)

assert source.streams_call_count == 1


def test_check_builds_a_named_state_delegating_stream(mocker):
Expand Down Expand Up @@ -1330,7 +1451,7 @@ def test_streams_by_name_enables_cache_for_checked_parent_without_mutating_manif
assert source.resolved_manifest == manifest_before


def test_check_still_fails_when_a_dynamic_stream_cannot_be_built(mocker):
def test_static_check_succeeds_when_a_dynamic_stream_cannot_be_built(mocker):
manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT)
manifest["dynamic_streams"] = [manifest["dynamic_streams"][1]]
manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]}
Expand Down Expand Up @@ -1361,8 +1482,10 @@ def create_component(*args, **kwargs):
)
status = source.check(logger, source._config)

assert status.status == Status.FAILED
assert "Encountered an error while discovering streams" in status.message
assert status.status == Status.SUCCEEDED

with pytest.raises(ValueError, match="bad dynamic stream"):
source.discover(logger=source.logger, config=source._config)


def _stream_group_manifest(streams: list[dict], check_stream_name: str) -> dict:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6650,14 +6650,16 @@ def test_get_partition_router(stream_factory, expected_type):


@pytest.mark.parametrize(
"build_path",
"build_path, expect_dynamic_stream_configs",
[
pytest.param("streams", id="streams"),
pytest.param("_streams_by_name", id="streams_by_name"),
pytest.param("_first_streams", id="first_streams"),
pytest.param("streams", True, id="streams"),
pytest.param("_streams_by_name", False, id="streams_by_name"),
pytest.param("_first_streams", True, id="first_streams"),
],
)
def test_api_budget_is_set_before_dynamic_streams_evaluated(build_path):
def test_api_budget_is_set_before_dynamic_streams_evaluated(
build_path, expect_dynamic_stream_configs
):
"""Verify all stream-building paths set the API budget before dynamic discovery."""
source = ConcurrentDeclarativeSource(
source_config=_MANIFEST, config=_CONFIG, catalog=None, state=None
Expand Down Expand Up @@ -6702,10 +6704,13 @@ def tracking_dynamic_stream_configs(*args, **kwargs):
source._first_streams(1)

assert "set_api_budget" in call_order, "set_api_budget was never called"
assert "dynamic_stream_configs" in call_order, "dynamic_stream_configs was never called"
assert call_order.index("set_api_budget") < call_order.index("dynamic_stream_configs"), (
f"set_api_budget must be called before dynamic_stream_configs, but call order was: {call_order}"
)
if expect_dynamic_stream_configs:
assert "dynamic_stream_configs" in call_order, "dynamic_stream_configs was never called"
assert call_order.index("set_api_budget") < call_order.index("dynamic_stream_configs"), (
f"set_api_budget must be called before dynamic_stream_configs, but call order was: {call_order}"
)
else:
assert "dynamic_stream_configs" not in call_order


@pytest.mark.parametrize("count", [0, -1])
Expand Down
Loading