diff --git a/dlt/sources/rest_api/config_setup.py b/dlt/sources/rest_api/config_setup.py index 3ce81f3aa7..bf62c6c4f7 100644 --- a/dlt/sources/rest_api/config_setup.py +++ b/dlt/sources/rest_api/config_setup.py @@ -66,6 +66,7 @@ Endpoint, EndpointResource, AuthType, + PaginatorType, ) @@ -105,6 +106,7 @@ def register_paginator( "Your custom paginator has to be a subclass of BasePaginator" ) PAGINATOR_MAP[paginator_name] = paginator_class + add_value_to_literal(PaginatorType, paginator_name) def get_paginator_class(paginator_name: str) -> Type[BasePaginator]: diff --git a/tests/sources/rest_api/configurations/test_custom_auth_config.py b/tests/sources/rest_api/configurations/test_custom_auth_config.py index 132bd67e88..52cdb95735 100644 --- a/tests/sources/rest_api/configurations/test_custom_auth_config.py +++ b/tests/sources/rest_api/configurations/test_custom_auth_config.py @@ -78,7 +78,7 @@ class NotAuthConfigBase: ) assert e.match("Invalid auth: NotAuthConfigBase.") - def test_validate_config_raises_no_error(self, custom_auth_config: AuthConfig) -> None: + def test_valid_config_raises_no_error(self, custom_auth_config: AuthConfig) -> None: rest_api.config_setup.register_auth("custom_oauth_2", CustomOAuth2) valid_config: RESTAPIConfig = { diff --git a/tests/sources/rest_api/configurations/test_custom_paginator_config.py b/tests/sources/rest_api/configurations/test_custom_paginator_config.py index f8ac060218..975ab10176 100644 --- a/tests/sources/rest_api/configurations/test_custom_paginator_config.py +++ b/tests/sources/rest_api/configurations/test_custom_paginator_config.py @@ -4,7 +4,7 @@ from dlt.sources import rest_api from dlt.sources.helpers.rest_client.paginators import JSONLinkPaginator -from dlt.sources.rest_api.typing import PaginatorConfig +from dlt.sources.rest_api.typing import PaginatorConfig, RESTAPIConfig class CustomPaginator(JSONLinkPaginator): @@ -67,3 +67,13 @@ class NotAPaginator: with pytest.raises(ValueError) as e: rest_api.config_setup.register_paginator("not_a_paginator", NotAPaginator) # type: ignore[arg-type] assert e.match("Invalid paginator: NotAPaginator.") + + def test_test_valid_config_raises_no_error(self, custom_paginator_config) -> None: + rest_api.config_setup.register_paginator("custom_paginator", CustomPaginator) + + valid_config: RESTAPIConfig = { + "client": {"base_url": "https://example.com", "paginator": custom_paginator_config}, + "resources": ["test"], + } + + rest_api.rest_api_source(valid_config)