Convert list_all_prs into iterator

pyproject.toml

@@ -42,6 +42,7 @@ mypy = "^1.4.1"
 pylint = "^2.17.5"
 pytest = "^7.4.0"
 pytest-cov = "^4.1.0"
+python-dotenv = "^1.0.1"
 PyYAML = "^6.0.1"
 toml = "^0.10.2"
 types-PyYAML = "^6.0.12.11"

simple_ado/__init__.py

@@ -6,7 +6,7 @@
 """ADO API wrapper."""
 
 import logging
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Tuple
 import urllib.parse
 
 import requests
@@ -22,7 +22,7 @@ from simple_ado.governance import ADOGovernanceClient
 from simple_ado.http_client import ADOHTTPClient, ADOResponse
 from simple_ado.pipelines import ADOPipelineClient
 from simple_ado.pools import ADOPoolsClient
-from simple_ado.pull_requests import ADOPullRequestClient
+from simple_ado.pull_requests import ADOPullRequestClient, ADOPullRequestStatus
 from simple_ado.security import ADOSecurityClient
 from simple_ado.user import ADOUserClient
 from simple_ado.wiki import ADOWikiClient
@@ -184,12 +184,15 @@ class ADOClient:
         branch_name: Optional[str] = None,
         project_id: str,
         repository_id: str,
-    ) -> ADOResponse:
+        top: int | None = None,
+        pr_status: ADOPullRequestStatus | None = None,
+    ) -> Iterator[Any]:
         """Get the pull requests for a branch from ADO.
 
         :param Optional[str] branch_name: The name of the branch to fetch the pull requests for.
         :param project_id: The ID of the project
         :param str repository_id: The ID for the repository
+        :param top: How many PRs to retrieve
 
         :returns: The ADO Response with the pull request data
         """
@@ -197,16 +200,24 @@ class ADOClient:
         self.log.debug("Fetching PRs")
 
         offset = 0
-        all_prs: List[Any] = []
 
         while True:
-
             request_url = (
                 self.http_client.api_endpoint(project_id=project_id)
                 + f"/git/repositories/{repository_id}/pullRequests?"
             )
 
-            request_url += f"$top=100&$skip={offset}"
+            parameters = {"$skip": offset}
+
+            if top:
+                parameters["$top"] = top
+
+            if pr_status:
+                parameters["searchCriteria.status"] = pr_status.value
+
+            encoded_parameters = urllib.parse.urlencode(parameters)
+
+            request_url += encoded_parameters
 
             if branch_name is not None:
                 request_url += f"&sourceRefName={_canonicalize_branch_name(branch_name)}"
@@ -221,11 +232,9 @@ class ADOClient:
             if len(extracted) == 0:
                 break
 
-            all_prs.extend(extracted)
+            yield from extracted
 
-            offset += 100
-
-        return all_prs
+            offset += len(extracted)
 
     def custom_get(
         self,

simple_ado/pull_requests.py

@@ -5,6 +5,7 @@
 
 """ADO Pull Request API wrapper."""
 
+import enum
 import logging
 from typing import Any, Dict, List, Optional
 
@@ -25,6 +26,16 @@ from simple_ado.http_client import ADOHTTPClient, ADOResponse, ADOThread
 from simple_ado.models import PatchOperation, AddOperation, DeleteOperation, PropertyValue
 
 
+class ADOPullRequestStatus(enum.Enum):
+    """Possible values of pull request states."""
+
+    ABANDONED: str = "abandoned"
+    ACTIVE: str = "active"
+    ALL: str = "all"
+    COMPLETED: str = "completed"
+    NOT_SET: str = "notSet"
+
+
 class ADOPullRequestClient(ADOBaseClient):
     """Wrapper class around the ADO Pull Request APIs.
 

tests/__init__.py

@@ -7,6 +7,9 @@ import os
 from typing import Any, Dict, List, Optional
 
 import toml
+from dotenv import load_dotenv
+
+load_dotenv()
 
 
 class TestDetails:

tests/test_simple_ado.py

@@ -8,6 +8,7 @@
 # pylint: disable=line-too-long,too-many-public-methods
 
 from collections import defaultdict
+import datetime
 import os
 import sys
 import unittest
@@ -341,3 +342,18 @@ class LibraryTests(unittest.TestCase):
             )
 
             assert full_definition is not None
+
+    def test_list_prs(self):
+        """Test list PRs diff."""
+        count = 0
+        one_month_ago = datetime.datetime.now() - datetime.timedelta(days=28)
+        for pr in self.client.list_all_pull_requests(
+            project_id=self.test_config.project_id,
+            repository_id=self.test_config.repository_id,
+            pr_status=simple_ado.pull_requests.ADOPullRequestStatus.COMPLETED,
+        ):
+            closed_date = datetime.datetime.strptime(pr["closedDate"][:-2], "%Y-%m-%dT%H:%M:%S.%f")
+            if closed_date < one_month_ago:
+                break
+            count += 1
+        assert count > 0