refactor(policy-check): Iterator conversion and unit tests added to the policy-check

Author: isivaselvan

examples/policy_check.py

@@ -35,7 +35,6 @@ def main():
         action="store_true",
         help="Get logs for the specified policy check",
     )
-    parser.add_argument("--page", type=int, default=1)
     parser.add_argument("--page-size", type=int, default=20)
     args = parser.parse_args()
 
@@ -50,21 +49,19 @@ def main():
     _print_header(f"Listing policy checks for run: {args.run_id}")
 
     options = PolicyCheckListOptions(
-        page_number=args.page,
         page_size=args.page_size,
     )
 
     try:
-        pc_list = client.policy_checks.list(args.run_id, options)
+        pc_list = list(client.policy_checks.list(args.run_id, options))
 
-        print(f"Total policy checks: {pc_list.total_count}")
-        print(f"Page {pc_list.current_page} of {pc_list.total_pages}")
+        print(f"Total policy checks fetched: {len(pc_list)}")
         print()
 
-        if not pc_list.items:
+        if not pc_list:
             print("No policy checks found for this run.")
         else:
-            for pc in pc_list.items:
+            for pc in pc_list:
                 print(f"- ID: {pc.id}")
                 print(f"Status: {pc.status}")
                 print(f"Scope: {pc.scope}")

src/pytfe/models/__init__.py

@@ -106,7 +106,6 @@
     PolicyActions,
     PolicyCheck,
     PolicyCheckIncludeOpt,
-    PolicyCheckList,
     PolicyCheckListOptions,
     PolicyPermissions,
     PolicyResult,
@@ -611,7 +610,6 @@
     "PolicyResult",
     "PolicyStatusTimestamps",
     "PolicyCheckListOptions",
-    "PolicyCheckList",
     # Policy Evaluation
     "PolicyAttachable",
     "PolicyEvaluation",
@@ -680,4 +678,3 @@
 
 # Rebuild models with forward references after all models are loaded
 PolicyCheck.model_rebuild()
-PolicyCheckList.model_rebuild()

src/pytfe/models/policy_check.py

@@ -106,7 +106,6 @@ class PolicyCheckListOptions(BaseModel):
     model_config = ConfigDict(populate_by_name=True, validate_by_name=True)
 
     include: list[PolicyCheckIncludeOpt] | None = Field(None, alias="include")
-    page_number: int | None = Field(None, alias="page[number]")
     page_size: int | None = Field(None, alias="page[size]")
 
 

src/pytfe/resources/policy_check.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Iterator
 
 from ..errors import (
     InvalidPolicyCheckIDError,
     InvalidRunIDError,
 )
 from ..models.policy_check import (
     PolicyCheck,
-    PolicyCheckList,
     PolicyCheckListOptions,
     PolicyStatus,
 )
@@ -24,35 +24,19 @@ class PolicyChecks(_Service):
 
     def list(
         self, run_id: str, options: PolicyCheckListOptions | None = None
-    ) -> PolicyCheckList:
+    ) -> Iterator[PolicyCheck]:
         """List all policy checks of the given run."""
         if not valid_string_id(run_id):
             raise InvalidRunIDError()
         params = (
             options.model_dump(by_alias=True, exclude_none=True) if options else None
         )
-        r = self.t.request(
-            "GET",
-            f"/api/v2/runs/{run_id}/policy-checks",
-            params=params,
-        )
-        jd = r.json()
-        items = []
-        meta = jd.get("meta", {})
-        pagination = meta.get("pagination", {})
-        for d in jd.get("data", []):
-            attrs = d.get("attributes", {})
-            attrs["id"] = d.get("id")
-            attrs["run"] = d.get("relationships", {}).get("run", {})
-            items.append(PolicyCheck.model_validate(attrs))
-        return PolicyCheckList(
-            items=items,
-            current_page=pagination.get("current-page"),
-            total_pages=pagination.get("total-pages"),
-            prev_page=pagination.get("prev-page"),
-            next_page=pagination.get("next-page"),
-            total_count=pagination.get("total-count"),
-        )
+        path = f"/api/v2/runs/{run_id}/policy-checks"
+        for item in self._list(path, params=params):
+            attrs = item.get("attributes", {})
+            attrs["id"] = item.get("id")
+            attrs["run"] = item.get("relationships", {}).get("run", {}).get("data")
+            yield PolicyCheck.model_validate(attrs)
 
     def read(self, policy_check_id: str) -> PolicyCheck:
         """Read a policy check by its ID."""
@@ -66,7 +50,7 @@ def read(self, policy_check_id: str) -> PolicyCheck:
         d = jd.get("data", {})
         attrs = d.get("attributes", {})
         attrs["id"] = d.get("id")
-        attrs["run"] = d.get("relationships", {}).get("run", {})
+        attrs["run"] = d.get("relationships", {}).get("run", {}).get("data")
         return PolicyCheck.model_validate(attrs)
 
     def override(self, policy_check_id: str) -> PolicyCheck:
@@ -81,7 +65,7 @@ def override(self, policy_check_id: str) -> PolicyCheck:
         d = jd.get("data", {})
         attrs = d.get("attributes", {})
         attrs["id"] = d.get("id")
-        attrs["run"] = d.get("relationships", {}).get("run", {})
+        attrs["run"] = d.get("relationships", {}).get("run", {}).get("data")
         return PolicyCheck.model_validate(attrs)
 
     def logs(self, policy_check_id: str) -> str: