Bug 1857600: Update pulse handler branch filtering

tests/etl/test_pulse_handler.py

@@ -0,0 +1,100 @@
+from unittest import mock
+
+from django.test import TestCase
+
+from treeherder.etl.taskcluster_pulse.handler import ignore_task
+from treeherder.model.models import Repository, RepositoryGroup
+
+
+@mock.patch("treeherder.etl.taskcluster_pulse.handler.projects_to_ingest", None)
+class TestPulseHandler(TestCase):
+    def setUp(self):
+        self.group = RepositoryGroup.objects.create(name="test-group")
+
+        # Existing mobile repo with master branch
+        self.ref_browser = Repository.objects.create(
+            name="reference-browser",
+            repository_group=self.group,
+            url="https://github.com/mozilla-mobile/reference-browser",
+            branch="master",
+            tc_root_url="https://firefox-ci-tc.services.mozilla.com",
+        )
+
+        # New repo (application-services) with main branch
+        self.app_services = Repository.objects.create(
+            name="application-services",
+            repository_group=self.group,
+            url="https://github.com/mozilla/application-services",
+            branch="main",
+            tc_root_url="https://firefox-ci-tc.services.mozilla.com",
+        )
+
+    def test_ignore_task_reference_browser_master(self):
+        """Test that reference-browser on master branch"""
+        task = {
+            "payload": {
+                "env": {
+                    "MOBILE_BASE_REPOSITORY": "https://github.com/mozilla-mobile/reference-browser",
+                    "MOBILE_HEAD_REPOSITORY": "https://github.com/mozilla-mobile/reference-browser",
+                    "MOBILE_HEAD_REF": "master",
+                }
+            }
+        }
+        # Should return False (do not ignore)
+        self.assertFalse(ignore_task(task, "task1", "root_url", "reference-browser"))
+
+    def test_ignore_task_reference_browser_feature(self):
+        """Test that reference-browser on feature branch is ignored"""
+        task = {
+            "payload": {
+                "env": {
+                    "MOBILE_BASE_REPOSITORY": "https://github.com/mozilla-mobile/reference-browser",
+                    "MOBILE_HEAD_REPOSITORY": "https://github.com/mozilla-mobile/reference-browser",
+                    "MOBILE_HEAD_REF": "feature-branch",
+                }
+            }
+        }
+        # Should return True (ignore)
+        self.assertTrue(ignore_task(task, "task2", "root_url", "reference-browser"))
+
+    def test_ignore_task_app_services_main(self):
+        """Test that application-services on main branch is not ignored"""
+        # Mocking decision task fallback
+        with mock.patch("treeherder.etl.taskcluster_pulse.handler.taskcluster.Queue") as mock_queue:
+            mock_client = mock.Mock()
+            mock_queue.return_value = mock_client
+
+            mock_client.task.return_value = {
+                "metadata": {
+                    "source": ["assume:repo:github.com/mozilla/application-services:branch:main"]
+                },
+                "routes": [],
+            }
+
+            task = {
+                "taskGroupId": "group1",
+                "payload": {
+                    "env": {}  # No MOBILE_ env vars
+                },
+            }
+
+            # Should return False (do not ignore)
+            self.assertFalse(ignore_task(task, "task3", "root_url", "application-services"))
+
+    def test_ignore_task_app_services_feature(self):
+        """Test that application-services on feature branch is ignored"""
+        with mock.patch("treeherder.etl.taskcluster_pulse.handler.taskcluster.Queue") as mock_queue:
+            mock_client = mock.Mock()
+            mock_queue.return_value = mock_client
+
+            mock_client.task.return_value = {
+                "metadata": {
+                    "source": ["assume:repo:github.com/mozilla/application-services:branch:feature"]
+                },
+                "routes": [],
+            }
+
+            task = {"taskGroupId": "group1", "payload": {"env": {}}}
+
+            # Should return True (ignore)
+            self.assertTrue(ignore_task(task, "task4", "root_url", "application-services"))

treeherder/etl/taskcluster_pulse/handler.py

@@ -12,6 +12,7 @@
 
 from treeherder.etl.schema import get_json_schema
 from treeherder.etl.taskcluster_pulse.parse_route import parse_route
+from treeherder.model.models import Repository
 
 env = environ.Env()
 logger = logging.getLogger(__name__)
@@ -111,55 +112,47 @@ def ignore_task(task, task_id, root_url, project):
         logger.debug("Ignoring tasks not matching PROJECTS_TO_INGEST (Task id: %s)", task_id)
         return True
 
-    mobile_repos = (
-        "reference-browser",
-        "mozilla-vpn-client",
-        "mozilla-vpn-client-release",
-    )
-    if project in mobile_repos:
+    repo = Repository.objects.filter(name=project).first()
+    if repo and repo.branch:
+        # Default to ignore, unless we find a matching branch
+        ignore = True
+
+        # Find branch in env vars
+        # Check if *any* env var ending in _HEAD_REF matches the repo branch
         envs = task["payload"].get("env", {})
-        if envs.get("MOBILE_BASE_REPOSITORY"):
-            try:
-                base_repo = envs["MOBILE_BASE_REPOSITORY"].rsplit("/", 1)[1]
-                if base_repo in mobile_repos:
-                    # Ignore tasks that are associated to a pull request
-                    if envs["MOBILE_BASE_REPOSITORY"] != envs["MOBILE_HEAD_REPOSITORY"]:
-                        logger.debug(
-                            "Task: %s belong to a pull request OR branch which we ignore.", task_id
-                        )
-                        ignore = True
-                    # Bug 1587542 - Temporary change to ignore Github tasks not associated to 'master'
-                    if envs["MOBILE_HEAD_REF"] not in (
-                        "refs/heads/master",
-                        "master",
-                        "refs/heads/main",
-                        "main",
-                    ):
-                        logger.info("Task: %s is not for the `master` branch.", task_id)
-                        ignore = True
-            except KeyError:
-                pass
-        else:
-            # The decision task is the ultimate source for determining this information
+        branch_found_in_env = False
+        for key, value in envs.items():
+            if key.endswith("_HEAD_REF"):
+                logger.debug(f"Checking env {key}={value} against repo.branch={repo.branch}")
+                branch_found_in_env = True
+                if value == repo.branch:
+                    ignore = False
+                    logger.info(f"Task {task_id} matched branch {repo.branch} via env var {key}")
+                    break
+
+        # If not found in env, check decision task (fallback)
+        if ignore and not branch_found_in_env:
+            # The decision task is the best source for determining this information
             queue = taskcluster.Queue({"rootUrl": root_url})
             decision_task = queue.task(task["taskGroupId"])
-            scopes = decision_task["metadata"].get("source")
-            ignore = True
+
+            # Check scopes
+            scopes = decision_task["metadata"].get("source", [])
             for scope in scopes:
-                # e.g. assume:repo:github.com/mozilla-mobile/fenix:branch:master
-                if scope.find("branch:master") != -1 or scope.find("branch:main") != -1:
+                if f"branch:{repo.branch}" in scope:
                     ignore = False
                     break
 
-            # This handles nightly tasks
-            # e.g. index.mobile.v2.fenix.branch.master.latest.taskgraph.decision-nightly
-            for route in decision_task["routes"]:
-                if route.find("master") != -1 or route.find("main") != -1:
-                    ignore = False
-                    break
+            # Check routes
+            if ignore:
+                for route in decision_task["routes"]:
+                    if repo.branch in route:  # Simple substring check as before
+                        ignore = False
+                        break
 
-    if ignore:
-        logger.debug(f"Task to be ignored ({task_id})")
+        if ignore:
+            logger.debug(f"Task {task_id} ignored: branch does not match '{repo.branch}'")
+            return True
 
     return ignore