From faab4c652c4943fc18c792995180bf59dbd5c7bc Mon Sep 17 00:00:00 2001
From: Summer Ji <summerji@google.com>
Date: Tue, 13 Jul 2021 09:05:20 -0700
Subject: [PATCH] feat: support map handle for DIREGAPIC Pagination (#1052)

* add implementation and test

* remove console.log

* add type for tuple

* Add detail resource check
---
 src/paginationCalls/pageDescriptor.ts | 11 +++-
 test/unit/pagedIteration.ts           | 85 ++++++++++++++++++++++++++-
 2 files changed, 93 insertions(+), 3 deletions(-)

diff --git a/src/paginationCalls/pageDescriptor.ts b/src/paginationCalls/pageDescriptor.ts
index 83367f7af..242e4d130 100644
--- a/src/paginationCalls/pageDescriptor.ts
+++ b/src/paginationCalls/pageDescriptor.ts
@@ -143,7 +143,7 @@ export class PageDescriptor implements Descriptor {
     const asyncIterable = {
       [Symbol.asyncIterator]() {
         let nextPageRequest: RequestType | null | undefined = request;
-        const cache: {}[] = [];
+        const cache: Array<ResponseType | [string, ResponseType]> = [];
         return {
           async next() {
             if (cache.length > 0) {
@@ -159,7 +159,14 @@ export class PageDescriptor implements Descriptor {
                 nextPageRequest!,
                 options
               )) as ResultTuple;
-              cache.push(...(result as ResponseType[]));
+              // For pagination response with protobuf map type, use tuple as representation.
+              if (result && !Array.isArray(result)) {
+                for (const [key, value] of Object.entries(result)) {
+                  cache.push([key, value as ResponseType]);
+                }
+              } else {
+                cache.push(...(result as ResponseType[]));
+              }
               if (cache.length === 0) {
                 ++attempts;
                 if (attempts > maxAttemptsEmptyResponse) {
diff --git a/test/unit/pagedIteration.ts b/test/unit/pagedIteration.ts
index 922e51d6e..bdaefa6d1 100644
--- a/test/unit/pagedIteration.ts
+++ b/test/unit/pagedIteration.ts
@@ -23,7 +23,7 @@ import * as sinon from 'sinon';
 import {PassThrough} from 'stream';
 import * as streamEvents from 'stream-events';
 import {PageDescriptor} from '../../src/paginationCalls/pageDescriptor';
-import {APICallback, GaxCall} from '../../src/apitypes';
+import {APICallback, GaxCall, RequestType} from '../../src/apitypes';
 import {describe, it, beforeEach} from 'mocha';
 
 import * as util from './utils';
@@ -473,3 +473,86 @@ describe('paged iteration', () => {
     });
   });
 });
+
+describe('REGAPIC Pagination', () => {
+  const pageSize = 3;
+  const pagesToStream = 5;
+  const descriptor = new PageDescriptor('pageToken', 'nextPageToken', 'items');
+  const retryOptions = util.createRetryOptions(0, 0, 0, 0, 0, 0, 100);
+  const createOptions = {
+    settings: {retry: retryOptions},
+    descriptor,
+  };
+  const response: RequestType = {
+    'regions/us-central1': {
+      warning: {
+        code: 'NO_RESULTS_ON_PAGE',
+        message:
+          "There are no results for scope 'regions/us-central1' on this page.",
+      },
+    },
+    'regions/us-east1': {
+      addresses: [
+        {
+          id: '5011754511478056813',
+          creationTimestamp: '2021-05-28T23:04:50.044-07:00',
+          name: 'test-address-0',
+        },
+        {
+          id: '1036412484008568684',
+          creationTimestamp: '2021-05-28T23:04:51.044-07:00',
+          name: 'test-address-1',
+        },
+      ],
+    },
+  };
+
+  function func(
+    request: {pageToken?: number},
+    metadata: {},
+    options: {},
+    callback: APICallback
+  ) {
+    const pageToken = request.pageToken || 0;
+    if (pageToken >= pageSize * pagesToStream) {
+      callback(null, {items: {}});
+    } else {
+      callback(null, {items: response, nextPageToken: pageToken + pageSize});
+    }
+  }
+
+  describe('use async tuple iterator', () => {
+    const spy = sinon.spy(func);
+    let apiCall: GaxCall;
+    beforeEach(() => {
+      apiCall = util.createApiCall(spy, createOptions);
+    });
+
+    async function iterableChecker(iterable: AsyncIterable<{} | undefined>) {
+      let counter = 0;
+      const resources = [];
+      for await (const resource of iterable) {
+        counter++;
+        resources.push(resource);
+        if (counter === 10) {
+          break;
+        }
+      }
+      return resources;
+    }
+
+    it('return an tuple iterable, count to 10', async () => {
+      const settings = new gax.CallSettings(
+        (createOptions && createOptions.settings) || {}
+      );
+      const resources = await iterableChecker(
+        descriptor.asyncIterate(apiCall, {}, settings)
+      );
+      assert.strictEqual(resources.length, 10);
+      for await (const [key, value] of resources as [string, ResponseType]) {
+        assert.ok(response.hasOwnProperty(key));
+        assert.strictEqual(value, response[key]);
+      }
+    });
+  });
+});