refactor(collections): use builder classes for collection request builders (issue #3272)

Signed-off-by: Sukuna0007Abhi <[email protected]>

Author: Sukuna0007Abhi

augur/tasks/start_tasks.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 import logging
+from abc import ABC, abstractmethod
+from typing import Any, List
 import os
 #from celery.result import AsyncResult
 from celery import group, chain
@@ -157,84 +159,133 @@ def non_repo_domain_tasks(self):
     tasks.apply_async()
 
 
-def build_primary_repo_collect_request(session, logger, enabled_phase_names, days_until_collect_again = 15):
-    #Add all required tasks to a list and pass it to the CollectionRequest
-    primary_enabled_phases = []
-    primary_gitlab_enabled_phases = []
+class CollectionRequestBuilder(ABC):
+    """Base class for collection request builders.
 
-    #Primary jobs
-    if prelim_phase.__name__ in enabled_phase_names:
-        primary_enabled_phases.append(prelim_phase)
+    Enforces the `build(session, logger, enabled_phase_names, days_until_collect_again)`
+    signature so all builders accept the same arguments.
+    """
 
-    primary_enabled_phases.append(primary_repo_collect_phase)
-    primary_gitlab_enabled_phases.append(primary_repo_collect_phase_gitlab)
+    @abstractmethod
+    def build(self, session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int) -> "CollectionRequest":
+        """Build and return a `CollectionRequest` configured for this builder."""
+        pass
 
-    #task success is scheduled no matter what the config says.
-    def core_task_success_util_gen(repo_git, full_collection):
-        return core_task_success_util.si(repo_git)
 
-    primary_enabled_phases.append(core_task_success_util_gen)
-    primary_gitlab_enabled_phases.append(core_task_success_util_gen)
+class PrimaryCollectionRequestBuilder(CollectionRequestBuilder):
+    def build(self, session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 15) -> "CollectionRequest":
+        """Return a `CollectionRequest` for primary collections.
 
-    primary_request = CollectionRequest("core",primary_enabled_phases,max_repo=40, days_until_collect_again=days_until_collect_again, gitlab_phases=primary_gitlab_enabled_phases)
-    primary_request.get_valid_repos(session)
-    return primary_request
+        Parameters mirror the runtime API used across task builders so callers can
+        pass `enabled_phase_names` and an interval `days_until_collect_again`.
+        """
+        # Build enabled phases
+        primary_enabled_phases: List = []
+        primary_gitlab_enabled_phases: List = []
 
-def build_secondary_repo_collect_request(session, logger, enabled_phase_names, days_until_collect_again = 1):
-    #Deal with secondary collection
-    secondary_enabled_phases = []
+        # Primary jobs
+        if prelim_phase.__name__ in enabled_phase_names:
+            primary_enabled_phases.append(prelim_phase)
 
-    if prelim_phase.__name__ in enabled_phase_names:
-        secondary_enabled_phases.append(prelim_phase_secondary)
+        primary_enabled_phases.append(primary_repo_collect_phase)
+        primary_gitlab_enabled_phases.append(primary_repo_collect_phase_gitlab)
 
+        # Add success task
+        def core_task_success_util_gen(repo_git, full_collection):
+            return core_task_success_util.si(repo_git)
 
-    secondary_enabled_phases.append(secondary_repo_collect_phase)
+        primary_enabled_phases.append(core_task_success_util_gen)
+        primary_gitlab_enabled_phases.append(core_task_success_util_gen)
 
-    def secondary_task_success_util_gen(repo_git, full_collection):
-        return secondary_task_success_util.si(repo_git)
+        primary_request = CollectionRequest("core",primary_enabled_phases,max_repo=40, days_until_collect_again=days_until_collect_again, gitlab_phases=primary_gitlab_enabled_phases)
+        primary_request.get_valid_repos(session)
+        return primary_request
 
-    secondary_enabled_phases.append(secondary_task_success_util_gen)
-    
-    request = CollectionRequest("secondary",secondary_enabled_phases,max_repo=60, days_until_collect_again=days_until_collect_again)
 
-    request.get_valid_repos(session)
-    return request
+class SecondaryCollectionRequestBuilder(CollectionRequestBuilder):
+    def build(self, session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 1) -> "CollectionRequest":
+        """Return a `CollectionRequest` for secondary collections."""
+        # Secondary collection
+        secondary_enabled_phases: List = []
+
+        if prelim_phase.__name__ in enabled_phase_names:
+            secondary_enabled_phases.append(prelim_phase_secondary)
+
+
+        secondary_enabled_phases.append(secondary_repo_collect_phase)
+
+        def secondary_task_success_util_gen(repo_git, full_collection):
+            return secondary_task_success_util.si(repo_git)
+
+        secondary_enabled_phases.append(secondary_task_success_util_gen)
+        
+        request = CollectionRequest("secondary",secondary_enabled_phases,max_repo=60, days_until_collect_again=days_until_collect_again)
+
+        request.get_valid_repos(session)
+        return request
+
+
+class FacadeCollectionRequestBuilder(CollectionRequestBuilder):
+    def build(self, session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 10) -> "CollectionRequest":
+        """Return a `CollectionRequest` for facade (git) collections."""
+        # Facade collection
+        facade_enabled_phases: List = []
+
+        facade_enabled_phases.append(facade_phase)
+
+        def facade_task_success_util_gen(repo_git, full_collection):
+            return facade_task_success_util.si(repo_git)
+
+        facade_enabled_phases.append(facade_task_success_util_gen)
+
+        def facade_task_update_weight_util_gen(repo_git, full_collection):
+            return git_update_commit_count_weight.si(repo_git)
+
+        facade_enabled_phases.append(facade_task_update_weight_util_gen)
+
+        request = CollectionRequest("facade",facade_enabled_phases,max_repo=30, days_until_collect_again=days_until_collect_again)
+
+        request.get_valid_repos(session)
+        return request
+
 
+class MLCollectionRequestBuilder(CollectionRequestBuilder):
+    def build(self, session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 40) -> "CollectionRequest":
+        """Return a `CollectionRequest` for machine-learning related collections."""
+        ml_enabled_phases: List = []
 
-def build_facade_repo_collect_request(session, logger, enabled_phase_names, days_until_collect_again = 10):
-    #Deal with facade collection
-    facade_enabled_phases = []
+        ml_enabled_phases.append(machine_learning_phase)
 
-    facade_enabled_phases.append(facade_phase)
+        def ml_task_success_util_gen(repo_git, full_collection):
+            return ml_task_success_util.si(repo_git)
 
-    def facade_task_success_util_gen(repo_git, full_collection):
-        return facade_task_success_util.si(repo_git)
+        ml_enabled_phases.append(ml_task_success_util_gen)
 
-    facade_enabled_phases.append(facade_task_success_util_gen)
+        request = CollectionRequest("ml",ml_enabled_phases,max_repo=5, days_until_collect_again=days_until_collect_again)
+        request.get_valid_repos(session)
+        return request
 
-    def facade_task_update_weight_util_gen(repo_git, full_collection):
-        return git_update_commit_count_weight.si(repo_git)
 
-    facade_enabled_phases.append(facade_task_update_weight_util_gen)
+# Backwards-compatible thin wrappers (keep API stable)
 
-    request = CollectionRequest("facade",facade_enabled_phases,max_repo=30, days_until_collect_again=days_until_collect_again)
+def build_primary_repo_collect_request(session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 15) -> "CollectionRequest":
+    """Thin wrapper around `PrimaryCollectionRequestBuilder.build` to preserve API compatibility."""
+    return PrimaryCollectionRequestBuilder().build(session, logger, enabled_phase_names, days_until_collect_again)
 
-    request.get_valid_repos(session)
-    return request
 
-def build_ml_repo_collect_request(session, logger, enabled_phase_names, days_until_collect_again = 40):
-    ml_enabled_phases = []
+def build_secondary_repo_collect_request(session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 1) -> "CollectionRequest":
+    """Thin wrapper around `SecondaryCollectionRequestBuilder.build` to preserve API compatibility."""
+    return SecondaryCollectionRequestBuilder().build(session, logger, enabled_phase_names, days_until_collect_again)
 
-    ml_enabled_phases.append(machine_learning_phase)
 
-    def ml_task_success_util_gen(repo_git, full_collection):
-        return ml_task_success_util.si(repo_git)
+def build_facade_repo_collect_request(session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 10) -> "CollectionRequest":
+    """Thin wrapper around `FacadeCollectionRequestBuilder.build` to preserve API compatibility."""
+    return FacadeCollectionRequestBuilder().build(session, logger, enabled_phase_names, days_until_collect_again)
 
-    ml_enabled_phases.append(ml_task_success_util_gen)
 
-    request = CollectionRequest("ml",ml_enabled_phases,max_repo=5, days_until_collect_again=days_until_collect_again)
-    request.get_valid_repos(session)
-    return request
+def build_ml_repo_collect_request(session: Any, logger: logging.Logger, enabled_phase_names: List[str], days_until_collect_again: int = 40) -> "CollectionRequest":
+    """Thin wrapper around `MLCollectionRequestBuilder.build` to preserve API compatibility."""
+    return MLCollectionRequestBuilder().build(session, logger, enabled_phase_names, days_until_collect_again)
 
 @celery.task(bind=True)
 def augur_collection_monitor(self):     
@@ -246,7 +297,7 @@ def augur_collection_monitor(self):
     logger.info("Checking for repos to collect")
 
 
-    #Get list of enabled phases 
+    # Get enabled phases
     enabled_phase_names = get_enabled_phase_names_from_config(engine, logger)
 
     enabled_collection_hooks = []

tests/test_tasks/test_collection_builders.py

@@ -0,0 +1,101 @@
+import inspect
+
+from augur.tasks.start_tasks import (
+    CollectionRequestBuilder,
+    PrimaryCollectionRequestBuilder,
+    SecondaryCollectionRequestBuilder,
+    FacadeCollectionRequestBuilder,
+    MLCollectionRequestBuilder,
+    build_primary_repo_collect_request,
+    build_secondary_repo_collect_request,
+    build_facade_repo_collect_request,
+    build_ml_repo_collect_request,
+)
+
+
+def test_builders_are_subclasses():
+    assert issubclass(PrimaryCollectionRequestBuilder, CollectionRequestBuilder)
+    assert issubclass(SecondaryCollectionRequestBuilder, CollectionRequestBuilder)
+    assert issubclass(FacadeCollectionRequestBuilder, CollectionRequestBuilder)
+    assert issubclass(MLCollectionRequestBuilder, CollectionRequestBuilder)
+
+
+def test_builders_have_build_method():
+    for cls in (
+        PrimaryCollectionRequestBuilder,
+        SecondaryCollectionRequestBuilder,
+        FacadeCollectionRequestBuilder,
+        MLCollectionRequestBuilder,
+    ):
+        assert hasattr(cls, "build")
+        assert inspect.isfunction(cls.build) or inspect.ismethod(cls.build)
+
+
+def test_wrapper_functions_exist():
+    assert callable(build_primary_repo_collect_request)
+    assert callable(build_secondary_repo_collect_request)
+    assert callable(build_facade_repo_collect_request)
+    assert callable(build_ml_repo_collect_request)
+
+
+def test_build_method_signature_contains_expected_params():
+    # Ensure the builders follow the expected signature so callers can pass enabled phase names
+    for cls in (
+        PrimaryCollectionRequestBuilder,
+        SecondaryCollectionRequestBuilder,
+        FacadeCollectionRequestBuilder,
+        MLCollectionRequestBuilder,
+    ):
+        sig = inspect.signature(cls.build)
+        params = list(sig.parameters.keys())
+        assert "enabled_phase_names" in params
+        assert "days_until_collect_again" in params
+
+
+def test_wrapper_signatures_match_build():
+    # Wrapper functions should accept the same parameters as the builders' build method
+    builders_and_wrappers = [
+        (PrimaryCollectionRequestBuilder, build_primary_repo_collect_request),
+        (SecondaryCollectionRequestBuilder, build_secondary_repo_collect_request),
+        (FacadeCollectionRequestBuilder, build_facade_repo_collect_request),
+        (MLCollectionRequestBuilder, build_ml_repo_collect_request),
+    ]
+
+    for builder_cls, wrapper in builders_and_wrappers:
+        builder_sig = inspect.signature(builder_cls.build)
+        wrapper_sig = inspect.signature(wrapper)
+        # Compare parameter names ignoring 'self' on the bound method
+        builder_params = [p for p in list(builder_sig.parameters.keys()) if p != 'self']
+        wrapper_params = list(wrapper_sig.parameters.keys())
+        assert builder_params == wrapper_params
+
+
+def test_build_returns_collection_request(monkeypatch):
+    import logging
+    from augur.tasks.start_tasks import PrimaryCollectionRequestBuilder
+    from augur.tasks.util.collection_util import CollectionRequest
+
+    # Avoid DB interaction by stubbing out get_valid_repos
+    monkeypatch.setattr(CollectionRequest, "get_valid_repos", lambda self, session: None)
+
+    builder = PrimaryCollectionRequestBuilder()
+    result = builder.build(session=None, logger=logging.getLogger("test"), enabled_phase_names=[], days_until_collect_again=1)
+
+    assert isinstance(result, CollectionRequest)
+
+
+def test_wrapper_is_passthrough(monkeypatch):
+    import logging
+    from augur.tasks.start_tasks import PrimaryCollectionRequestBuilder, build_primary_repo_collect_request
+    from augur.tasks.util.collection_util import CollectionRequest
+
+    monkeypatch.setattr(CollectionRequest, "get_valid_repos", lambda self, session: None)
+
+    b = PrimaryCollectionRequestBuilder()
+    res_builder = b.build(session=None, logger=logging.getLogger("test"), enabled_phase_names=["a"], days_until_collect_again=2)
+    res_wrapper = build_primary_repo_collect_request(session=None, logger=logging.getLogger("test"), enabled_phase_names=["a"], days_until_collect_again=2)
+
+    # Basic sanity checks to ensure wrapper returns what the builder builds
+    assert isinstance(res_wrapper, CollectionRequest)
+    assert res_builder.name == res_wrapper.name
+    assert res_builder.max_repo == res_wrapper.max_repo