Move callGAPIItems and add unit tests

src/gam.py

@@ -995,36 +995,6 @@ def callGAPIpages(service, function, items='items',
         sys.stderr.flush()
       return all_items
 
-def callGAPIitems(service, function, items='items',
-                  throw_reasons=None, retry_reasons=None,
-                  **kwargs):
-  """Gets a single page of items from a Google service function that is paged.
-
-  Args:
-    service: A Google service object for the desired API.
-    function: String, The name of a service request method to execute.
-    items: String, the name of the resulting "items" field within the service
-        method's response object.
-    soft_errors: Bool, If True, writes non-fatal errors to stderr.
-    throw_reasons: A list of Google HTTP error reason strings indicating the
-        errors generated by this request should be re-thrown. All other HTTP
-        errors are consumed.
-    retry_reasons: A list of Google HTTP error reason strings indicating which
-        error should be retried, using exponential backoff techniques, when the
-        error reason is encountered.
-
-  Returns:
-    The list of items in the first page of a response.
-  """
-  results = gapi.call(service,
-                     function,
-                     throw_reasons=throw_reasons,
-                     retry_reasons=retry_reasons,
-                     **kwargs)
-  if results:
-    return results.get(items, [])
-  return []
-
 def getAPIVersion(api):
   version = API_VER_MAPPING.get(api, 'v1')
   if api in ['directory', 'reports', 'datatransfer']:
@@ -1484,9 +1454,10 @@ def showReport():
     while True:
       try:
         if fullDataRequired is not None:
-          warnings = callGAPIitems(rep.userUsageReport(), 'get', 'warnings',
-                                   throw_reasons=[gapi.errors.ErrorReason.INVALID],
-                                   date=tryDate, userKey=userKey, customerId=customerId, orgUnitID=orgUnitId, fields='warnings')
+          warnings = gapi.get_first_page(
+              rep.userUsageReport(), 'get', 'warnings',
+              throw_reasons=[gapi.errors.ErrorReason.INVALID],
+              date=tryDate, userKey=userKey, customerId=customerId, orgUnitID=orgUnitId, fields='warnings')
           fullData, tryDate = _checkFullDataAvailable(warnings, tryDate, fullDataRequired)
           if fullData < 0:
             print('No user report available.')
@@ -1526,9 +1497,10 @@ def showReport():
     while True:
       try:
         if fullDataRequired is not None:
-          warnings = callGAPIitems(rep.customerUsageReports(), 'get', 'warnings',
-                                   throw_reasons=[gapi.errors.ErrorReason.INVALID],
-                                   customerId=customerId, date=tryDate, fields='warnings')
+          warnings = gapi.get_first_page(
+              rep.customerUsageReports(), 'get', 'warnings',
+              throw_reasons=[gapi.errors.ErrorReason.INVALID],
+              customerId=customerId, date=tryDate, fields='warnings')
           fullData, tryDate = _checkFullDataAvailable(warnings, tryDate, fullDataRequired)
           if fullData < 0:
             print('No customer report available.')
@@ -1972,8 +1944,9 @@ def doGetCustomerInfo():
   # of current primary being added, not customer create date.
   # We should also get all domains and use oldest date
   oldest = datetime.datetime.strptime(customer_info['customerCreationTime'], '%Y-%m-%dT%H:%M:%S.%fZ')
-  domains = callGAPIitems(cd.domains(), 'list', 'domains',
-                          customer=GC_Values[GC_CUSTOMER_ID], fields='domains(creationTime)')
+  domains = gapi.get_first_page(cd.domains(), 'list', 'domains',
+                                customer=GC_Values[GC_CUSTOMER_ID],
+                                fields='domains(creationTime)')
   for domain in domains:
     domain_creation = datetime.datetime.fromtimestamp(int(domain['creationTime'])/1000)
     if domain_creation < oldest:
@@ -7818,8 +7791,9 @@ def convertGCPFolderNameToID(parent, crm2):
   # crm2.folders() is broken requiring pageToken, etc in body, not URL.
   # for now just use callGAPI and if user has that many folders they'll
   # just need to be specific.
-  folders = callGAPIitems(crm2.folders(), 'search', items='folders',
-                          body={'pageSize': 1000, 'query': 'displayName="%s"' % parent})
+  folders = gapi.get_first_page(
+      crm2.folders(), 'search', items='folders',
+      body={'pageSize': 1000, 'query': 'displayName="%s"' % parent})
   if not folders:
     controlflow.system_error_exit(1, 'ERROR: No folder found matching displayName=%s' % parent)
   if len(folders) > 1:
@@ -10550,7 +10524,7 @@ def doSiteVerifyShow():
 
 def doGetSiteVerifications():
   verif = buildGAPIObject('siteVerification')
-  sites = callGAPIitems(verif.webResource(), 'list', 'items')
+  sites = gapi.get_first_page(verif.webResource(), 'list', 'items')
   if sites:
     for site in sites:
       print('Site: %s' % site['site']['identifier'])
@@ -10763,7 +10737,7 @@ def doGetOrgInfo(name=None, return_attrib=None):
 def doGetASPs(users):
   cd = buildGAPIObject('directory')
   for user in users:
-    asps = callGAPIitems(cd.asps(), 'list', 'items', userKey=user)
+    asps = gapi.get_first_page(cd.asps(), 'list', 'items', userKey=user)
     if asps:
       print('Application-Specific Passwords for %s' % user)
       for asp in asps:
@@ -10789,7 +10763,7 @@ def doDelASP(users):
     codeIds = codeIdList.replace(',', ' ').split()
   for user in users:
     if allCodeIds:
-      asps = callGAPIitems(cd.asps(), 'list', 'items', userKey=user, fields='items/codeId')
+      asps = gapi.get_first_page(cd.asps(), 'list', 'items', userKey=user, fields='items/codeId')
       codeIds = [asp['codeId'] for asp in asps]
     for codeId in codeIds:
       gapi.call(cd.asps(), 'delete', userKey=user, codeId=codeId)
@@ -10814,7 +10788,7 @@ def doGetBackupCodes(users):
   cd = buildGAPIObject('directory')
   for user in users:
     try:
-      codes = callGAPIitems(cd.verificationCodes(), 'list', 'items', throw_reasons=[gapi.errors.ErrorReason.INVALID_ARGUMENT, gapi.errors.ErrorReason.INVALID], userKey=user)
+      codes = gapi.get_first_page(cd.verificationCodes(), 'list', 'items', throw_reasons=[gapi.errors.ErrorReason.INVALID_ARGUMENT, gapi.errors.ErrorReason.INVALID], userKey=user)
     except (gapi.errors.GapiInvalidArgumentError, gapi.errors.GapiInvalidError):
       codes = []
     printBackupCodes(user, codes)
@@ -10823,7 +10797,7 @@ def doGenBackupCodes(users):
   cd = buildGAPIObject('directory')
   for user in users:
     gapi.call(cd.verificationCodes(), 'generate', userKey=user)
-    codes = callGAPIitems(cd.verificationCodes(), 'list', 'items', userKey=user)
+    codes = gapi.get_first_page(cd.verificationCodes(), 'list', 'items', userKey=user)
     printBackupCodes(user, codes)
 
 def doDelBackupCodes(users):
@@ -10909,9 +10883,9 @@ def _showToken(token):
                             throw_reasons=[gapi.errors.ErrorReason.NOT_FOUND, gapi.errors.ErrorReason.USER_NOT_FOUND, gapi.errors.ErrorReason.RESOURCE_NOT_FOUND],
                             userKey=user, clientId=clientId, fields=fields)]
       else:
-        results = callGAPIitems(cd.tokens(), 'list', 'items',
-                                throw_reasons=[gapi.errors.ErrorReason.USER_NOT_FOUND],
-                                userKey=user, fields='items({0})'.format(fields))
+        results = gapi.get_first_page(cd.tokens(), 'list', 'items',
+                                      throw_reasons=[gapi.errors.ErrorReason.USER_NOT_FOUND],
+                                      userKey=user, fields='items({0})'.format(fields))
       jcount = len(results)
       if not csvFormat:
         print('User: {0}, Access Tokens ({1}/{2})'.format(user, i, count))
@@ -10937,7 +10911,7 @@ def doDeprovUser(users):
   cd = buildGAPIObject('directory')
   for user in users:
     print('Getting Application Specific Passwords for %s' % user)
-    asps = callGAPIitems(cd.asps(), 'list', 'items', userKey=user, fields='items/codeId')
+    asps = gapi.get_first_page(cd.asps(), 'list', 'items', userKey=user, fields='items/codeId')
     jcount = len(asps)
     if jcount > 0:
       j = 0
@@ -10953,7 +10927,7 @@ def doDeprovUser(users):
     except gapi.errors.GapiInvalidError:
       print('No 2SV Backup Codes')
     print('Getting tokens for %s...' % user)
-    tokens = callGAPIitems(cd.tokens(), 'list', 'items', userKey=user, fields='items/clientId')
+    tokens = gapi.get_first_page(cd.tokens(), 'list', 'items', userKey=user, fields='items/clientId')
     jcount = len(tokens)
     if jcount > 0:
       j = 0

src/gapi/__init__.py

@@ -138,6 +138,41 @@ def call(service,
       controlflow.system_error_exit(4, str(e))
 
 
+def get_first_page(service,
+                   function,
+                   items='items',
+                   throw_reasons=None,
+                   retry_reasons=None,
+                   **kwargs):
+  """Gets a single page of items from a Google service function that is paged.
+
+  Args:
+    service: A Google service object for the desired API.
+    function: String, The name of a service request method to execute.
+    items: String, the name of the resulting "items" field within the service
+      method's response object.
+    throw_reasons: A list of Google HTTP error reason strings indicating the
+      errors generated by this request should be re-thrown. All other HTTP
+      errors are consumed.
+    retry_reasons: A list of Google HTTP error reason strings indicating which
+      error should be retried, using exponential backoff techniques, when the
+      error reason is encountered.
+    **kwargs: Additional params to pass to the request method.
+
+  Returns:
+    The list of items in the first page of a response.
+  """
+  results = call(
+      service,
+      function,
+      throw_reasons=throw_reasons,
+      retry_reasons=retry_reasons,
+      **kwargs)
+  if results:
+    return results.get(items, [])
+  return []
+
+
 # TODO: Make this private once all execution related items that use this method
 # have been brought into this file
 def handle_oauth_token_error(e, soft_errors):

src/gapi/__init___test.py

@@ -72,14 +72,14 @@ def test_create_http_sets_cache_timeout(self):
     self.assertEqual(http.timeout, 1234)
 
 
-class CallTest(unittest.TestCase):
+class GapiTest(unittest.TestCase):
 
   def setUp(self):
     SetGlobalVariables()
     self.mock_service = MagicMock()
     self.mock_method_name = 'mock_method'
     self.mock_method = getattr(self.mock_service, self.mock_method_name)
-    super(CallTest, self).setUp()
+    super(GapiTest, self).setUp()
 
   def test_call_returns_basic_200_response(self):
     response = gapi.call(self.mock_service, self.mock_method_name)
@@ -236,6 +236,38 @@ def test_call_retries_requests_with_backoff_on_servernotfounderror(
     # Make sure a backoff technique was used for retry.
     self.assertEqual(mock_wait_on_failure.call_count, 1)
 
+  def test_get_first_page_calls_correct_service_function(self):
+    pass
+
+  def test_get_first_page_returns_one_page(self):
+    fake_response = {'items': [{}, {}, {}]}
+    self.mock_method.return_value.execute.return_value = fake_response
+    page = gapi.get_first_page(self.mock_service, self.mock_method_name)
+    self.assertEqual(page, fake_response['items'])
+
+  def test_get_first_page_non_standard_page_field_name(self):
+    field_name = 'things'
+    fake_response = {field_name: [{}, {}, {}]}
+    self.mock_method.return_value.execute.return_value = fake_response
+    page = gapi.get_first_page(
+        self.mock_service, self.mock_method_name, items=field_name)
+    self.assertEqual(page, fake_response[field_name])
+
+  def test_get_first_page_passes_additional_kwargs_to_service(self):
+    gapi.get_first_page(
+        self.mock_service, self.mock_method_name, my_param_1=1, my_param_2=2)
+    self.mock_method.assert_called_once()
+    method_kwargs = self.mock_method.call_args[1]
+    self.assertEqual(1, method_kwargs.get('my_param_1'))
+    self.assertEqual(2, method_kwargs.get('my_param_2'))
+
+  def test_get_first_page_returns_empty_list_when_no_items_returned(self):
+    non_items_response = {'noItemsInThisResponse': {}}
+    self.mock_method.return_value.execute.return_value = non_items_response
+    page = gapi.get_first_page(self.mock_service, self.mock_method_name)
+    self.assertIsInstance(page, list)
+    self.assertEqual(0, len(page))
+
 
 if __name__ == '__main__':
   unittest.main()