From 95df8d0f16f44ce10fbfff046fd4f7879f3b8d3e Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Wed, 24 Jan 2024 15:59:59 -0500
Subject: [PATCH] fix: ensure /github/branches endpoint returns all branches

This endoint currently only fetches up to 100 branches from Github as
this is the maximum number of objects you can get from a single query to
the GraphQL endpoint.

But `mozilla-vpn-client` has well over 100 branches, and the `release`
branches were being excluded from the query. Fix this by using
pagination until all branches have been returned.

Issue: #1392
---
 api/pyproject.toml                 |  1 +
 api/src/shipit_api/admin/github.py | 53 ++++++++++++++++++++----------
 api/tests/conftest.py              |  7 ++++
 api/tests/test_github.py           | 52 +++++++++++++++++++++++++++++
 4 files changed, 96 insertions(+), 17 deletions(-)
 create mode 100644 api/tests/test_github.py

diff --git a/api/pyproject.toml b/api/pyproject.toml
index 97af8484a..a6cf9b7b2 100644
--- a/api/pyproject.toml
+++ b/api/pyproject.toml
@@ -93,6 +93,7 @@ responses = "*"
 pytest = "*"
 pytest-asyncio = "*"
 pytest-cov = "*"
+pytest-mock = "*"
 black = "*"
 check-manifest = "*"
 isort = "*"
diff --git a/api/src/shipit_api/admin/github.py b/api/src/shipit_api/admin/github.py
index 8a4198b43..17024f33b 100644
--- a/api/src/shipit_api/admin/github.py
+++ b/api/src/shipit_api/admin/github.py
@@ -114,27 +114,46 @@ def ref_to_commit(owner, repo, ref):
 
 
 def list_github_branches(owner, repo, limit=100):
-    query = """
-    {
-      repository(name: "%(repo)s", owner: "%(owner)s") {
-        refs(first: %(limit)s, refPrefix: "refs/heads/") {
-          nodes {
-            name
-            target {
-              ... on Commit {
-                committedDate
+    page = {
+        "hasNextPage": True,
+        "endCursor": None,
+    }
+
+    branches = []
+    while page["hasNextPage"]:
+        after = ""
+        if page["endCursor"]:
+            after = f", after: \"{page['endCursor']}\""
+
+        query = """
+        {
+          repository(name: "%(repo)s", owner: "%(owner)s") {
+            refs(first: %(limit)s%(after)s, refPrefix: "refs/heads/") {
+              nodes {
+                name
+                target {
+                  ... on Commit {
+                    committedDate
+                  }
+                }
+              }
+              pageInfo {
+                endCursor
+                hasNextPage
               }
             }
           }
         }
-      }
-    }
-    """ % dict(
-        owner=owner, repo=repo, limit=limit
-    )
-    content = query_api(query)
-    nodes = content["data"]["repository"]["refs"]["nodes"]
-    return [{"committer_date": node["target"]["committedDate"], "name": node["name"]} for node in nodes]
+        """ % dict(
+            owner=owner, repo=repo, limit=limit, after=after
+        )
+        content = query_api(query)
+        page = content["data"]["repository"]["refs"]["pageInfo"]
+
+        nodes = content["data"]["repository"]["refs"]["nodes"]
+        branches.extend([{"committer_date": node["target"]["committedDate"], "name": node["name"]} for node in nodes])
+
+    return branches
 
 
 def list_github_commits(owner, repo, branch, limit=10):
diff --git a/api/tests/conftest.py b/api/tests/conftest.py
index cf0c38808..57de9f4e2 100644
--- a/api/tests/conftest.py
+++ b/api/tests/conftest.py
@@ -6,11 +6,18 @@
 import os
 
 import pytest
+import responses as rsps
 
 import backend_common.testing
 from shipit_api.common.models import XPI, XPIRelease
 
 
+@pytest.fixture
+def responses():
+    with rsps.RequestsMock() as req_mock:
+        yield req_mock
+
+
 @pytest.fixture(scope="session")
 def app():
     """Load shipit_api in test mode"""
diff --git a/api/tests/test_github.py b/api/tests/test_github.py
new file mode 100644
index 000000000..b0495aebb
--- /dev/null
+++ b/api/tests/test_github.py
@@ -0,0 +1,52 @@
+from pprint import pprint
+
+import pytest
+from flask import Flask
+
+from shipit_api.admin import github
+
+
+@pytest.fixture(autouse=True)
+def setup(monkeypatch, mocker):
+    app = Flask(__name__)
+    app.config["GITHUB_TOKEN"] = "token"
+    with app.app_context():
+        monkeypatch.setattr(github, "current_app", app)
+        mocker.patch("shipit_api.admin.github._require_auth", return_value=None)
+        yield
+
+
+def test_list_github_branches(responses):
+    rsp1 = responses.post(
+        "https://api.github.com/graphql",
+        json={
+            "data": {
+                "repository": {
+                    "refs": {
+                        "nodes": [{"name": "foo", "target": {"committedDate": "1"}}, {"name": "bar", "target": {"committedDate": "2"}}],
+                        "pageInfo": {"hasNextPage": True, "endCursor": "abc"},
+                    }
+                }
+            }
+        },
+    )
+    rsp2 = responses.post(
+        "https://api.github.com/graphql",
+        json={
+            "data": {
+                "repository": {
+                    "refs": {
+                        "nodes": [{"name": "baz", "target": {"committedDate": "3"}}],
+                        "pageInfo": {"hasNextPage": False, "endCursor": "def"},
+                    }
+                }
+            }
+        },
+    )
+    repos = github.list_github_branches("org", "repo", limit=2)
+    assert b"after" not in rsp1.calls[0].request.body
+    assert b"after: \\\"abc\\\"," in rsp2.calls[0].request.body
+
+    print("Dumping for copy/paste:")
+    pprint(repos)
+    assert repos == [{"committer_date": "1", "name": "foo"}, {"committer_date": "2", "name": "bar"}, {"committer_date": "3", "name": "baz"}]