firebase · hiranya911 · Mar 23, 2020 · Mar 21, 2020 · Mar 23, 2020 · Mar 23, 2020
diff --git a/firebase_admin/tenant_mgt.py b/firebase_admin/tenant_mgt.py
@@ -27,14 +27,19 @@
 
 
 _TENANT_MGT_ATTRIBUTE = '_tenant_mgt'
+_MAX_LIST_TENANTS_RESULTS = 100
 
 
 __all__ = [
+    'ListTenantsPage',
     'Tenant',
     'TenantNotFoundError',
 
+    'create_tenant',
     'delete_tenant',
     'get_tenant',
+    'list_tenants',
+    'update_tenant',
 ]
 
 TenantNotFoundError = _auth_utils.TenantNotFoundError
@@ -128,6 +133,34 @@ def delete_tenant(tenant_id, app=None):
     tenant_mgt_service.delete_tenant(tenant_id)
 
 
+def list_tenants(page_token=None, max_results=_MAX_LIST_TENANTS_RESULTS, app=None):
+    """Retrieves a page of tenants from a Firebase project.
+
+    The ``page_token`` argument governs the starting point of the page. The ``max_results``
+    argument governs the maximum number of tenants that may be included in the returned page.
+    This function never returns None. If there are no user accounts in the Firebase project, this
+    returns an empty page.
+
+    Args:
+        page_token: A non-empty page token string, which indicates the starting point of the page
+            (optional). Defaults to ``None``, which will retrieve the first page of users.
+        max_results: A positive integer indicating the maximum number of users to include in the
+            returned page (optional). Defaults to 100, which is also the maximum number allowed.
+        app: An App instance (optional).
+
+    Returns:
+        ListTenantsPage: A ListTenantsPage instance.
+
+    Raises:
+        ValueError: If ``max_results`` or ``page_token`` are invalid.
+        FirebaseError: If an error occurs while retrieving the user accounts.
+    """
+    tenant_mgt_service = _get_tenant_mgt_service(app)
+    def download(page_token, max_results):
+        return tenant_mgt_service.list_tenants(page_token, max_results)
+    return ListTenantsPage(download, page_token, max_results)
+
+
 def _get_tenant_mgt_service(app):
     return _utils.get_app_service(app, _TENANT_MGT_ATTRIBUTE, _TenantManagementService)
 
@@ -254,3 +287,106 @@ def delete_tenant(self, tenant_id):
             self.client.request('delete', '/tenants/{0}'.format(tenant_id))
         except requests.exceptions.RequestException as error:
             raise _auth_utils.handle_auth_backend_error(error)
+
+    def list_tenants(self, page_token=None, max_results=_MAX_LIST_TENANTS_RESULTS):
+        """Retrieves a batch of tenants."""
+        if page_token is not None:
+            if not isinstance(page_token, str) or not page_token:
+                raise ValueError('Page token must be a non-empty string.')
+        if not isinstance(max_results, int):
+            raise ValueError('Max results must be an integer.')
+        if max_results < 1 or max_results > _MAX_LIST_TENANTS_RESULTS:
+            raise ValueError(
+                'Max results must be a positive integer less than or equal to '
+                '{0}.'.format(_MAX_LIST_TENANTS_RESULTS))
+
+        payload = {'pageSize': max_results}
+        if page_token:
+            payload['pageToken'] = page_token
+        try:
+            return self.client.body('get', '/tenants', params=payload)
+        except requests.exceptions.RequestException as error:
+            raise _auth_utils.handle_auth_backend_error(error)
+
+
+class ListTenantsPage:
+    """Represents a page of tenants fetched from a Firebase project.
+
+    Provides methods for traversing tenants included in this page, as well as retrieving
+    subsequent pages of tenants. The iterator returned by ``iterate_all()`` can be used to iterate
+    through all tenants in the Firebase project starting from this page.
+    """
+
+    def __init__(self, download, page_token, max_results):
+        self._download = download
+        self._max_results = max_results
+        self._current = download(page_token, max_results)
+
+    @property
+    def tenants(self):
+        """A list of ``ExportedUserRecord`` instances available in this page."""
+        return [Tenant(data) for data in self._current.get('tenants', [])]
+
+    @property
+    def next_page_token(self):
+        """Page token string for the next page (empty string indicates no more pages)."""
+        return self._current.get('nextPageToken', '')
+
+    @property
+    def has_next_page(self):
+        """A boolean indicating whether more pages are available."""
+        return bool(self.next_page_token)
+
+    def get_next_page(self):
+        """Retrieves the next page of tenants, if available.
+
+        Returns:
+            ListTenantsPage: Next page of tenants, or None if this is the last page.
+        """
+        if self.has_next_page:
+            return ListTenantsPage(self._download, self.next_page_token, self._max_results)
+        return None
+
+    def iterate_all(self):
+        """Retrieves an iterator for tenants.
+
+        Returned iterator will iterate through all the tenants in the Firebase project
+        starting from this page. The iterator will never buffer more than one page of tenants
+        in memory at a time.
+
+        Returns:
+            iterator: An iterator of Tenant instances.
+        """
+        return _TenantIterator(self)
+
+
+class _TenantIterator:
+    """An iterator that allows iterating over tenants.
+
+    This implementation loads a page of tenants into memory, and iterates on them. When the whole
+    page has been traversed, it loads another page. This class never keeps more than one page
+    of entries in memory.
+    """
+
+    def __init__(self, current_page):
+        if not current_page:
+            raise ValueError('Current page must not be None.')
+        self._current_page = current_page
+        self._index = 0
+
+    def next(self):
+        if self._index == len(self._current_page.tenants):
+            if self._current_page.has_next_page:
+                self._current_page = self._current_page.get_next_page()
+                self._index = 0
+        if self._index < len(self._current_page.tenants):
+            result = self._current_page.tenants[self._index]
+            self._index += 1
+            return result
+        raise StopIteration
+
+    def __next__(self):
+        return self.next()
+
+    def __iter__(self):
+        return self
diff --git a/tests/test_tenant_mgt.py b/tests/test_tenant_mgt.py
@@ -15,6 +15,7 @@
 """Test cases for the firebase_admin.tenant_mgt module."""
 
 import json
+from urllib import parse
 
 import pytest
 
@@ -37,6 +38,38 @@
     }
 }"""
 
+LIST_TENANTS_RESPONSE = """{
+    "tenants": [
+        {
+            "name": "projects/mock-project-id/tenants/tenant0",
+            "displayName": "Test Tenant",
+            "allowPasswordSignup": true,
+            "enableEmailLinkSignin": true
+        },
+        {
+            "name": "projects/mock-project-id/tenants/tenant1",
+            "displayName": "Test Tenant",
+            "allowPasswordSignup": true,
+            "enableEmailLinkSignin": true
+        }
+    ]
+}"""
+
+LIST_TENANTS_RESPONSE_WITH_TOKEN = """{
+    "tenants": [
+        {
+            "name": "projects/mock-project-id/tenants/tenant0"
+        },
+        {
+            "name": "projects/mock-project-id/tenants/tenant1"
+        },
+        {
+            "name": "projects/mock-project-id/tenants/tenant2"
+        }
+    ],
+    "nextPageToken": "token"
+}"""
+
 INVALID_TENANT_IDS = [None, '', 0, 1, True, False, list(), tuple(), dict()]
 INVALID_BOOLEANS = ['', 1, 0, list(), tuple(), dict()]
 
@@ -309,8 +342,142 @@ def test_tenant_not_found(self, tenant_mgt_app):
         assert excinfo.value.cause is not None
 
 
-def _assert_tenant(tenant):
-    assert tenant.tenant_id == 'tenant-id'
+class TestListTenants:
+
+    @pytest.mark.parametrize('arg', [None, 'foo', list(), dict(), 0, -1, 101, False])
+    def test_invalid_max_results(self, tenant_mgt_app, arg):
+        with pytest.raises(ValueError):
+            tenant_mgt.list_tenants(max_results=arg, app=tenant_mgt_app)
+
+    @pytest.mark.parametrize('arg', ['', list(), dict(), 0, -1, True, False])
+    def test_invalid_page_token(self, tenant_mgt_app, arg):
+        with pytest.raises(ValueError):
+            tenant_mgt.list_tenants(page_token=arg, app=tenant_mgt_app)
+
+    def test_list_single_page(self, tenant_mgt_app):
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE)
+        page = tenant_mgt.list_tenants(app=tenant_mgt_app)
+        self._assert_tenants_page(page)
+        assert page.next_page_token == ''
+        assert page.has_next_page is False
+        assert page.get_next_page() is None
+        tenants = [tenant for tenant in page.iterate_all()]
+        assert len(tenants) == 2
+        self._assert_request(recorder)
+
+    def test_list_multiple_pages(self, tenant_mgt_app):
+        # Page 1
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE_WITH_TOKEN)
+        page = tenant_mgt.list_tenants(app=tenant_mgt_app)
+        assert len(page.tenants) == 3
+        assert page.next_page_token == 'token'
+        assert page.has_next_page is True
+        self._assert_request(recorder)
+
+        # Page 2 (also the last page)
+        response = {'tenants': [{'name': 'projects/mock-project-id/tenants/tenant3'}]}
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, json.dumps(response))
+        page = page.get_next_page()
+        assert len(page.tenants) == 1
+        assert page.next_page_token == ''
+        assert page.has_next_page is False
+        assert page.get_next_page() is None
+        self._assert_request(recorder, {'pageSize': '100', 'pageToken': 'token'})
+
+    def test_list_tenants_paged_iteration(self, tenant_mgt_app):
+        # Page 1
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE_WITH_TOKEN)
+        page = tenant_mgt.list_tenants(app=tenant_mgt_app)
+        iterator = page.iterate_all()
+        for index in range(3):
+            tenant = next(iterator)
+            assert tenant.tenant_id == 'tenant{0}'.format(index)
+        self._assert_request(recorder)
+
+        # Page 2 (also the last page)
+        response = {'tenants': [{'name': 'projects/mock-project-id/tenants/tenant3'}]}
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, json.dumps(response))
+        tenant = next(iterator)
+        assert tenant.tenant_id == 'tenant3'
+
+        with pytest.raises(StopIteration):
+            next(iterator)
+        self._assert_request(recorder, {'pageSize': '100', 'pageToken': 'token'})
+
+    def test_list_tenants_iterator_state(self, tenant_mgt_app):
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE)
+        page = tenant_mgt.list_tenants(app=tenant_mgt_app)
+
+        # Advance iterator.
+        iterator = page.iterate_all()
+        tenant = next(iterator)
+        assert tenant.tenant_id == 'tenant0'
+
+        # Iterator should resume from where left off.
+        tenant = next(iterator)
+        assert tenant.tenant_id == 'tenant1'
+
+        with pytest.raises(StopIteration):
+            next(iterator)
+        self._assert_request(recorder)
+
+    def test_list_tenants_stop_iteration(self, tenant_mgt_app):
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE)
+        page = tenant_mgt.list_tenants(app=tenant_mgt_app)
+        iterator = page.iterate_all()
+        tenants = [tenant for tenant in iterator]
+        assert len(tenants) == 2
+
+        with pytest.raises(StopIteration):
+            next(iterator)
+        self._assert_request(recorder)
+
+    def test_list_tenants_no_tenants_response(self, tenant_mgt_app):
+        response = {'tenants': []}
+        _instrument_tenant_mgt(tenant_mgt_app, 200, json.dumps(response))
+        page = tenant_mgt.list_tenants(app=tenant_mgt_app)
+        assert len(page.tenants) == 0
+        tenants = [tenant for tenant in page.iterate_all()]
+        assert len(tenants) == 0
+
+    def test_list_tenants_with_max_results(self, tenant_mgt_app):
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE)
+        page = tenant_mgt.list_tenants(max_results=50, app=tenant_mgt_app)
+        self._assert_tenants_page(page)
+        self._assert_request(recorder, {'pageSize' : '50'})
+
+    def test_list_tenants_with_all_args(self, tenant_mgt_app):
+        _, recorder = _instrument_tenant_mgt(tenant_mgt_app, 200, LIST_TENANTS_RESPONSE)
+        page = tenant_mgt.list_tenants(page_token='foo', max_results=50, app=tenant_mgt_app)
+        self._assert_tenants_page(page)
+        self._assert_request(recorder, {'pageToken' : 'foo', 'pageSize' : '50'})
+
+    def test_list_tenants_error(self, tenant_mgt_app):
+        _instrument_tenant_mgt(tenant_mgt_app, 500, '{"error":"test"}')
+        with pytest.raises(exceptions.InternalError) as excinfo:
+            tenant_mgt.list_tenants(app=tenant_mgt_app)
+        assert str(excinfo.value) == 'Unexpected error response: {"error":"test"}'
+
+    def _assert_tenants_page(self, page):
+        assert isinstance(page, tenant_mgt.ListTenantsPage)
+        assert len(page.tenants) == 2
+        for idx, tenant in enumerate(page.tenants):
+            _assert_tenant(tenant, 'tenant{0}'.format(idx))
+
+    def _assert_request(self, recorder, expected=None):
+        if expected is None:
+            expected = {'pageSize' : '100'}
+
+        assert len(recorder) == 1
+        req = recorder[0]
+        assert req.method == 'GET'
+        request = dict(parse.parse_qsl(parse.urlsplit(req.url).query))
+        assert request == expected
+
+
+def _assert_tenant(tenant, tenant_id='tenant-id'):
+    assert isinstance(tenant, tenant_mgt.Tenant)
+    assert tenant.tenant_id == tenant_id
     assert tenant.display_name == 'Test Tenant'
     assert tenant.allow_password_sign_up is True
     assert tenant.enable_email_link_sign_in is True