From f389d507e8054083d8a3962a02af9b2cfa247ef9 Mon Sep 17 00:00:00 2001
From: Brian Pandola <bpandola@gmail.com>
Date: Fri, 13 May 2022 19:40:29 -0700
Subject: [PATCH] Add Pagination for
 Organizations:ListOrganizationalUnitsForParent

---
 moto/organizations/models.py                  | 15 +++++++--------
 moto/organizations/responses.py               | 16 ++++++++++++----
 moto/organizations/utils.py                   |  7 +++++++
 .../test_organizations_boto3.py               | 19 +++++++++++++++++++
 4 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/moto/organizations/models.py b/moto/organizations/models.py
index 7140fdba7..98ced4794 100644
--- a/moto/organizations/models.py
+++ b/moto/organizations/models.py
@@ -434,15 +434,14 @@ class OrganizationsBackend(BaseBackend):
         ou = self.get_organizational_unit_by_id(kwargs["OrganizationalUnitId"])
         return ou.describe()
 
+    @paginate(pagination_model=PAGINATION_MODEL)
     def list_organizational_units_for_parent(self, **kwargs):
-        parent_id = self.validate_parent_id(kwargs["ParentId"])
-        return dict(
-            OrganizationalUnits=[
-                {"Id": ou.id, "Arn": ou.arn, "Name": ou.name}
-                for ou in self.ou
-                if ou.parent_id == parent_id
-            ]
-        )
+        parent_id = self.validate_parent_id(kwargs["parent_id"])
+        return [
+            {"Id": ou.id, "Arn": ou.arn, "Name": ou.name}
+            for ou in self.ou
+            if ou.parent_id == parent_id
+        ]
 
     def create_account(self, **kwargs):
         new_account = FakeAccount(self.org, **kwargs)
diff --git a/moto/organizations/responses.py b/moto/organizations/responses.py
index c731a321a..c0268e86b 100644
--- a/moto/organizations/responses.py
+++ b/moto/organizations/responses.py
@@ -51,11 +51,19 @@ class OrganizationsResponse(BaseResponse):
         )
 
     def list_organizational_units_for_parent(self):
-        return json.dumps(
-            self.organizations_backend.list_organizational_units_for_parent(
-                **self.request_params
-            )
+        max_results = self._get_int_param("MaxResults")
+        next_token = self._get_param("NextToken")
+        parent_id = self._get_param("ParentId")
+        (
+            ous,
+            next_token,
+        ) = self.organizations_backend.list_organizational_units_for_parent(
+            max_results=max_results, next_token=next_token, parent_id=parent_id
         )
+        response = {"OrganizationalUnits": ous}
+        if next_token:
+            response["NextToken"] = next_token
+        return json.dumps(response)
 
     def list_parents(self):
         return json.dumps(
diff --git a/moto/organizations/utils.py b/moto/organizations/utils.py
index 596f4d5be..95f00254b 100644
--- a/moto/organizations/utils.py
+++ b/moto/organizations/utils.py
@@ -41,6 +41,13 @@ PAGINATION_MODEL = {
         "result_key": "Accounts",
         "unique_attribute": "JoinedTimestamp",
     },
+    "list_organizational_units_for_parent": {
+        "input_token": "next_token",
+        "limit_key": "max_results",
+        "limit_default": 20,
+        "result_key": "OrganizationalUnits",
+        "unique_attribute": "Id",
+    },
 }
 
 
diff --git a/tests/test_organizations/test_organizations_boto3.py b/tests/test_organizations/test_organizations_boto3.py
index 8d02309f0..a37049dec 100644
--- a/tests/test_organizations/test_organizations_boto3.py
+++ b/tests/test_organizations/test_organizations_boto3.py
@@ -143,6 +143,25 @@ def test_list_organizational_units_for_parent():
         validate_organizational_unit(org, dict(OrganizationalUnit=ou))
 
 
+@mock_organizations
+def test_list_organizational_units_pagination():
+    client = boto3.client("organizations", region_name="us-east-1")
+    client.create_organization(FeatureSet="ALL")
+    root_id = client.list_roots()["Roots"][0]["Id"]
+    for i in range(20):
+        name = "ou" + str(i)
+        client.create_organizational_unit(ParentId=root_id, Name=name)
+    response = client.list_organizational_units_for_parent(ParentId=root_id)
+    response.should_not.have.key("NextToken")
+    len(response["OrganizationalUnits"]).should.be.greater_than_or_equal_to(i)
+
+    paginator = client.get_paginator("list_organizational_units_for_parent")
+    page_iterator = paginator.paginate(MaxResults=5, ParentId=root_id)
+    for page in page_iterator:
+        len(page["OrganizationalUnits"]).should.be.lower_than_or_equal_to(5)
+    page["OrganizationalUnits"][-1]["Name"].should.contain("19")
+
+
 @mock_organizations
 def test_list_organizational_units_for_parent_exception():
     client = boto3.client("organizations", region_name="us-east-1")