diff --git a/examples/policy.py b/examples/policy.py index 74352f4..4629cca 100644 --- a/examples/policy.py +++ b/examples/policy.py @@ -82,7 +82,6 @@ def main(): _print_header(f"Listing policies in organization: {args.org}") list_options = PolicyListOptions( - page_number=args.page, page_size=args.page_size, ) @@ -93,14 +92,10 @@ def main(): PolicyKind.SENTINEL if args.kind == "sentinel" else PolicyKind.OPA ) - policy_list = client.policies.list(args.org, list_options) - - print(f"Total policies: {policy_list.total_count}") - print(f"Page {policy_list.current_page} of {policy_list.total_pages}") - print() + policy_iter = client.policies.list(args.org, list_options) existing_policy = None - for policy in policy_list.items: + for policy in policy_iter: print( f"- {policy.id} | {policy.name} | kind={policy.kind} | enforcement={policy.enforcement_level}" ) diff --git a/src/pytfe/models/policy.py b/src/pytfe/models/policy.py index fb182d9..daf614f 100644 --- a/src/pytfe/models/policy.py +++ b/src/pytfe/models/policy.py @@ -38,7 +38,6 @@ class PolicyListOptions(BaseModel): search: str | None = Field(None, alias="search[name]") kind: PolicyKind | None = Field(None, alias="filter[kind]") - page_number: int | None = Field(None, alias="page[number]") page_size: int | None = Field(None, alias="page[size]") diff --git a/src/pytfe/resources/policy.py b/src/pytfe/resources/policy.py index fb30ca0..25153c0 100644 --- a/src/pytfe/resources/policy.py +++ b/src/pytfe/resources/policy.py @@ -1,5 +1,8 @@ from __future__ import annotations +from collections.abc import Iterator +from typing import Any + from ..errors import ( InvalidNameError, InvalidOrgError, @@ -11,7 +14,6 @@ from ..models.policy import ( Policy, PolicyCreateOptions, - PolicyList, PolicyListOptions, PolicyUpdateOptions, ) @@ -22,35 +24,28 @@ class Policies(_Service): def list( self, organization: str, options: PolicyListOptions | None = None - ) -> PolicyList: - """List all the policies of the given organization.""" + ) -> Iterator[Policy]: + """Iterate all the policies of the given organization.""" if not valid_string_id(organization): raise InvalidOrgError() - params = ( - options.model_dump(by_alias=True, exclude_none=True) if options else None - ) - r = self.t.request( - "GET", - f"/api/v2/organizations/{organization}/policies", - 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["organization"] = d.get("relationships", {}).get("organization", {}) - items.append(Policy.model_validate(attrs)) - return PolicyList( - 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/organizations/{organization}/policies" + params: dict[str, Any] = {} + + if options: + if getattr(options, "page_size", None): + params["page[size]"] = str(options.page_size) + + def _gen() -> Iterator[Policy]: + for item in self._list(path, params=params): + attrs = item.get("attributes", {}) + attrs["id"] = item.get("id") + attrs["organization"] = item.get("relationships", {}).get( + "organization", {} + ) + yield Policy.model_validate(attrs) + + return _gen() def create(self, organization: str, options: PolicyCreateOptions) -> Policy: """Create a new policy in the given organization.""" diff --git a/tests/units/test_policy.py b/tests/units/test_policy.py index af4790c..ea335f9 100644 --- a/tests/units/test_policy.py +++ b/tests/units/test_policy.py @@ -14,7 +14,6 @@ EnforcementLevel, Policy, PolicyCreateOptions, - PolicyList, PolicyUpdateOptions, ) from pytfe.models.policy_set import PolicyKind @@ -84,22 +83,16 @@ def test_list_policies_success_without_options( mock_response.json.return_value = mock_response_data mock_transport.request.return_value = mock_response - result = policies_service.list("org-123") + result_iter = policies_service.list("org-123") + items = list(result_iter) - mock_transport.request.assert_called_once_with( - "GET", "/api/v2/organizations/org-123/policies", params=None - ) + assert mock_transport.request.called - assert isinstance(result, PolicyList) - assert len(result.items) == 1 - assert result.items[0].id == "pol-123" - assert result.items[0].name == "test-policy" - assert result.items[0].kind == PolicyKind.SENTINEL - assert ( - result.items[0].enforcement_level == EnforcementLevel.ENFORCEMENT_ADVISORY - ) - assert result.current_page == 1 - assert result.total_count == 1 + assert len(items) == 1 + assert items[0].id == "pol-123" + assert items[0].name == "test-policy" + assert items[0].kind == PolicyKind.SENTINEL + assert items[0].enforcement_level == EnforcementLevel.ENFORCEMENT_ADVISORY def test_create_policy_validations(self, policies_service): """Test create method validations."""