handle next token when getting services for application

uktrade · Jan 13, 2025 · 4689235 · 4689235
1 parent 37f4181
commit 4689235
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 3 deletions.
diff --git a/dbt_platform_helper/utils/application.py b/dbt_platform_helper/utils/application.py
@@ -60,7 +60,7 @@ def __eq__(self, other):
         return str(self) == str(other)
 
 
-def load_application(app: str = None, default_session: Session = None) -> Application:
+def load_application(app: str | None = None, default_session: Session | None = None) -> Application:
     application = Application(app if app else get_application_name())
     current_session = default_session if default_session else get_aws_session_or_abort()
 
@@ -72,7 +72,7 @@ def load_application(app: str = None, default_session: Session = None) -> Applic
             WithDecryption=False,
         )
     except ssm_client.exceptions.ParameterNotFound:
-        raise ApplicationNotFoundException(app)
+        raise ApplicationNotFoundException(application.name)
 
     path = f"/copilot/applications/{application.name}/environments"
     secrets = get_ssm_secrets(app, None, current_session, path)
@@ -106,10 +106,19 @@ def is_environment_key(name):
         Recursive=False,
         WithDecryption=False,
     )
+    results = response["Parameters"]
+    while "NextToken" in response:
+        response = ssm_client.get_parameters_by_path(
+            Path=f"/copilot/applications/{application.name}/components",
+            Recursive=False,
+            WithDecryption=False,
+            NextToken=response["NextToken"],
+        )
+        results.extend(response["Parameters"])
 
     application.services = {
         svc["name"]: Service(svc["name"], svc["type"])
-        for svc in [json.loads(parameter["Value"]) for parameter in response["Parameters"]]
+        for svc in [json.loads(parameter["Value"]) for parameter in results]
     }
 
     return application

diff --git a/tests/platform_helper/utils/test_application.py b/tests/platform_helper/utils/test_application.py
@@ -254,3 +254,46 @@ def test_loading_an_application_in_a_different_account(
         )
 
         self.assertRaises(ApplicationNotFoundException, load_application, "sample")
+
+    def test_loading_an_application_services_with_token(
+        self, get_aws_session_or_abort, get_profile_name_from_account_id
+    ):
+        session = MagicMock(name="session-mock")
+        client = MagicMock(name="client-mock")
+        session.client.return_value = client
+
+        client.get_caller_identity.return_value = {"Account": "abc_123"}
+        client.get_parameters_by_path.side_effect = [
+            {
+                "Parameters": [
+                    {
+                        "Name": "/copilot/applications/another-test/environments/my_env",
+                        "Value": '{"name": "my_env", "accountID": "abc_123"}',
+                    }
+                ]
+            },
+            {
+                "Parameters": [
+                    {
+                        "Name": "/copilot/applications/another-test/components/web",
+                        "Value": '{"app": "demoddjango", "name": "web", "type": "Load Balanced Web Service"}',
+                    }
+                ],
+                "NextToken": "sometoken",
+            },
+            {
+                "Parameters": [
+                    {
+                        "Name": "/copilot/applications/another-test/components/web11",
+                        "Value": '{"app": "demoddjango", "name": "web11", "type": "Load Balanced Web Service"}',
+                    }
+                ]
+            },
+        ]
+        application = load_application(app="another-test", default_session=session)
+
+        self.assertEqual(len(application.services), 2)
+        self.assertEqual(application.services["web"].name, "web")
+        self.assertEqual(application.services["web"].kind, "Load Balanced Web Service")
+        self.assertEqual(application.services["web11"].name, "web11")
+        self.assertEqual(application.services["web11"].kind, "Load Balanced Web Service")