fix broken cli test

Netflix · Jan 19, 2021 · 8ec9a0c · 8ec9a0c
1 parent 4e29319
commit 8ec9a0c
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 23 deletions.
diff --git a/repokid/datasource/iam/__init__.py b/repokid/datasource/iam/__init__.py
@@ -14,6 +14,7 @@
 
 import copy
 import logging
+from typing import Dict
 from typing import Optional
 
 from cloudaux.aws.iam import get_account_authorization_details
@@ -31,13 +32,7 @@ class IAMDatasource(DatasourcePlugin[str, IAMEntry], Singleton):
     def __init__(self, config: Optional[RepokidConfig] = None):
         super().__init__(config=config)
 
-    def get(self, arn: str) -> IAMEntry:
-        result = self._data.get(arn)
-        if not result:
-            raise NotFoundError
-        return result
-
-    def seed(self, account_number: str) -> None:
+    def _fetch(self, account_number: str) -> Dict[str, IAMEntry]:
         logger.info("getting role data for account %s", account_number)
         conn = copy.deepcopy(self.config["connection_iam"])
         conn["account_number"] = account_number
@@ -49,7 +44,17 @@ def seed(self, account_number: str) -> None:
                 item["PolicyName"]: item["PolicyDocument"]
                 for item in data["RolePolicyList"]
             }
-        self._data.update(auth_details_by_id)
+        return auth_details_by_id
+
+    def get(self, arn: str) -> IAMEntry:
+        result = self._data.get(arn)
+        if not result:
+            raise NotFoundError
+        return result
+
+    def seed(self, account_number: str) -> None:
+        fetched_data = self._fetch(account_number)
+        self._data.update(fetched_data)
 
 
 class ConfigDatasource(DatasourcePlugin[str, IAMEntry], Singleton):

diff --git a/tests/test_commands.py b/tests/test_commands.py
@@ -273,10 +273,12 @@ class TestRepokidCLI(object):
     @patch("repokid.commands.role_cache.find_and_mark_inactive")
     @patch("repokid.commands.role_cache.RoleList.store")
     @patch("repokid.commands.role_cache.Role.gather_role_data")
-    @patch("repokid.datasource.iam.get_account_authorization_details")
+    @patch("repokid.commands.role_cache.AccessAdvisorDatasource")
+    @patch("repokid.datasource.iam.IAMDatasource._fetch")
     def test_repokid_update_role_cache(
         self,
-        mock_get_account_authorization_details,
+        mock_iam_datasource_fetch,
+        mock_access_advisor_datasource,
         mock_gather_role_data,
         mock_role_list_store,
         mock_find_and_mark_inactive,
@@ -300,8 +302,9 @@ def test_repokid_update_role_cache(
                 "PolicyDocument": ROLE_POLICIES["all_services_used"],
             }
         ]
+        role_data = {item["RoleId"]: item for item in role_data}
 
-        mock_get_account_authorization_details.side_effect = [role_data]
+        mock_iam_datasource_fetch.return_value = role_data
 
         config = {
             "aardvark_api_location": "",
@@ -323,18 +326,16 @@ def test_repokid_update_role_cache(
         assert mock_gather_role_data.call_count == 3
 
         # all roles active
-        assert mock_find_and_mark_inactive.mock_calls == [
-            call(
-                account_number,
-                RoleList(
-                    [
-                        Role(**ROLES[0]),
-                        Role(**ROLES[1]),
-                        Role(**ROLES[2]),
-                    ]
-                ),
-            )
-        ]
+        assert mock_find_and_mark_inactive.mock_calls[-1] == call(
+            account_number,
+            RoleList(
+                [
+                    Role(**ROLES[0]),
+                    Role(**ROLES[1]),
+                    Role(**ROLES[2]),
+                ]
+            ),
+        )
 
     @patch("tabview.view")
     @patch("repokid.commands.role.RoleList.from_ids")