fix: test for gql query 300 limit

googleapis · Nov 9, 2022 · 2377993 · 2377993
1 parent 0fc89a1
commit 2377993
Show file tree

Hide file tree

Showing 3 changed files with 83 additions and 3 deletions.
diff --git a/Datastore/src/EntityPageIterator.php b/Datastore/src/EntityPageIterator.php
@@ -61,6 +61,19 @@ public function current()
         $this->moreResultsType = isset($this->page['batch']['moreResults'])
             ? $this->page['batch']['moreResults']
             : null;
+        if ($this->config['resultLimit'] === 0 &&
+            isset($this->page['query']['limit'])
+        ) {
+            // limit is not set properly, so set the limit from parsed query
+            // which is returned by the server together with batch response
+            $limit = $this->page['query']['limit'];
+            if (is_array($limit) &&
+                isset($this->page['query']['limit']['value'])
+            ) {
+                $limit = $this->page['query']['limit']['value'];
+            }
+            $this->config['resultLimit'] = $limit;
+        }
 
         return $this->get($this->itemsPath, $this->page);
     }

diff --git a/Datastore/tests/System/QueryResultPaginationTest.php b/Datastore/tests/System/QueryResultPaginationTest.php
@@ -134,6 +134,19 @@ public function testGqlQueryPaginationByPage(DatastoreClient $client)
         $this->assertQueryPageCount(self::$expectedTotal, $client, $q);
     }
 
+    /**
+     * @dataProvider clientProvider
+     */
+    public function testGqlQueryPaginationWithLimit(DatastoreClient $client)
+    {
+        $testLimit = 310;
+        $q = $client->gqlQuery(
+            sprintf('SELECT * FROM `%s` LIMIT %s', self::$testKind, $testLimit),
+            ['allowLiterals' => true]
+        );
+        $this->assertQueryCount($testLimit, $client, $q);
+    }
+
     /**
      * @dataProvider clientProvider
      */

diff --git a/Datastore/tests/Unit/EntityPageIteratorTest.php b/Datastore/tests/Unit/EntityPageIteratorTest.php
@@ -40,6 +40,60 @@ function ($item) {
 
         $pagesArray = iterator_to_array($pages);
 
+        $this->assertOverPages($pages);
+    }
+
+    public function testCurrentSetsQueryIntegerLimits()
+    {
+        $pages = new EntityPageIterator(
+            function ($item) {
+                return $item;
+            },
+            [$this, 'theCall'],
+            [
+                'test' => 'call',
+                'query' => ['limit' => 1000],
+            ]
+        );
+
+        $pagesArray = iterator_to_array($pages);
+
+        $this->assertOverPages($pages, 1000);
+    }
+
+    public function testCurrentSetsQueryArrayLimits()
+    {
+
+        $pages = new EntityPageIterator(
+            function ($item) {
+                return $item;
+            },
+            [$this, 'theCall'],
+            [
+                'test' => 'call',
+                'query' => ['limit' => ['value' => 1100]],
+            ]
+        );
+
+        $this->assertOverPages($pages, 1100);
+    }
+
+    private function assertOverPages($pages, $expectedLimit = 0){
+        $pagesArray = iterator_to_array($pages);
+
+        $reflection = new \ReflectionClass($pages);
+        $configsProp = $reflection->getProperty('config');
+        $configsProp->setAccessible(true);
+        $configs = $configsProp->getValue($pages);
+
+        $this->assertIsArray($configs);
+        $this->assertArrayHasKey('resultLimit', $configs);
+        $this->assertEquals(
+            $expectedLimit,
+            $configs['resultLimit'],
+            "resultLimit assertion failed."
+        );
+
         $this->assertEquals(self::$moreResultsType, $pages->moreResultsType());
         $this->assertEquals(self::$items, $pagesArray[0]);
     }
@@ -48,9 +102,9 @@ public function theCall(array $options)
     {
         return [
             'batch' => [
-                'moreResults' => self::$moreResultsType
+                'moreResults' => self::$moreResultsType,
             ],
-            'items' => self::$items
-        ];
+            'items' => self::$items,
+        ] + $options;
     }
 }