diff --git a/swh/scheduler/backend.py b/swh/scheduler/backend.py --- a/swh/scheduler/backend.py +++ b/swh/scheduler/backend.py @@ -1,4 +1,4 @@ -# Copyright (C) 2015-2021 The Software Heritage developers +# Copyright (C) 2015-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information @@ -314,6 +314,7 @@ self, lister_id: Optional[UUID] = None, url: Optional[str] = None, + enabled_only: bool = True, limit: int = 1000, page_token: Optional[ListedOriginPageToken] = None, db=None, @@ -334,6 +335,9 @@ query_filters.append("url = %s") query_params.append(url) + if enabled_only: + query_filters.append("enabled = true") + if page_token is not None: query_filters.append("(lister_id, url) > %s") # the typeshed annotation for tuple() is too strict. diff --git a/swh/scheduler/interface.py b/swh/scheduler/interface.py --- a/swh/scheduler/interface.py +++ b/swh/scheduler/interface.py @@ -1,4 +1,4 @@ -# Copyright (C) 2015-2021 The Software Heritage developers +# Copyright (C) 2015-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information @@ -385,6 +385,7 @@ self, lister_id: Optional[UUID] = None, url: Optional[str] = None, + enabled_only: bool = True, limit: int = 1000, page_token: Optional[ListedOriginPageToken] = None, ) -> PaginatedListedOriginList: diff --git a/swh/scheduler/tests/conftest.py b/swh/scheduler/tests/conftest.py --- a/swh/scheduler/tests/conftest.py +++ b/swh/scheduler/tests/conftest.py @@ -1,4 +1,4 @@ -# Copyright (C) 2016-2021 The Software Heritage developers +# Copyright (C) 2016-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information @@ -63,6 +63,14 @@ return sum(listed_origins_by_type.values(), []) +@pytest.fixture +def listed_origins_with_non_enabled(listed_origins) -> List[ListedOrigin]: + """Return a (fixed) set of listed origins""" + for i, origin in enumerate(listed_origins): + origin.enabled = i % 2 == 0 + return listed_origins + + @pytest.fixture def storage(swh_storage): """An instance of in-memory storage that gets injected diff --git a/swh/scheduler/tests/test_journal_client.py b/swh/scheduler/tests/test_journal_client.py --- a/swh/scheduler/tests/test_journal_client.py +++ b/swh/scheduler/tests/test_journal_client.py @@ -1,4 +1,4 @@ -# Copyright (C) 2021 The Software Heritage developers +# Copyright (C) 2021-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information @@ -1003,7 +1003,7 @@ ) # Now check that the origin in question is disabled - actual_page = swh_scheduler.get_listed_origins(url="bar") + actual_page = swh_scheduler.get_listed_origins(url="bar", enabled_only=False) assert len(actual_page.results) == 1 assert actual_page.next_page_token is None diff --git a/swh/scheduler/tests/test_scheduler.py b/swh/scheduler/tests/test_scheduler.py --- a/swh/scheduler/tests/test_scheduler.py +++ b/swh/scheduler/tests/test_scheduler.py @@ -1,4 +1,4 @@ -# Copyright (C) 2017-2021 The Software Heritage developers +# Copyright (C) 2017-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information @@ -795,6 +795,27 @@ assert ret.next_page_token is None assert len(ret.results) == len(listed_origins) + def test_get_listed_origins_enabled_only( + self, swh_scheduler, listed_origins_with_non_enabled + ) -> None: + swh_scheduler.record_listed_origins(listed_origins_with_non_enabled) + ret = swh_scheduler.get_listed_origins( + enabled_only=True, limit=len(listed_origins_with_non_enabled) + 1 + ) + assert ret.next_page_token is None + assert len(ret.results) == len( + [lo for lo in listed_origins_with_non_enabled if lo.enabled] + ) + assert all([lo.enabled for lo in ret.results]) + + swh_scheduler.record_listed_origins(listed_origins_with_non_enabled) + ret = swh_scheduler.get_listed_origins( + enabled_only=False, limit=len(listed_origins_with_non_enabled) + 1 + ) + assert ret.next_page_token is None + assert len(ret.results) == len(listed_origins_with_non_enabled) + assert any([not lo.enabled for lo in ret.results]) + def _grab_next_visits_setup(self, swh_scheduler, listed_origins_by_type): """Basic origins setup for scheduling policy tests""" visit_type = next(iter(listed_origins_by_type))