From bb6e587a088fb36a36d9525626e4dcc4a215c57c Mon Sep 17 00:00:00 2001
From: Brent Shaffer <betterbrent@google.com>
Date: Tue, 18 Jun 2024 10:19:54 -0700
Subject: [PATCH] add tests for logger

---
 Bigtable/src/BigtableClient.php          |  2 +
 Bigtable/src/ChunkFormatter.php          |  3 +-
 Bigtable/src/Table.php                   |  5 +-
 Bigtable/tests/Unit/SmartRetriesTest.php | 59 +++++++++++++++++++++++-
 4 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/Bigtable/src/BigtableClient.php b/Bigtable/src/BigtableClient.php
index 5e519616bef5..5ec0aa7de518 100644
--- a/Bigtable/src/BigtableClient.php
+++ b/Bigtable/src/BigtableClient.php
@@ -23,6 +23,7 @@
 use Google\Auth\FetchAuthTokenInterface;
 use Google\Cloud\Bigtable\V2\BigtableClient as GapicClient;
 use Google\Cloud\Core\ArrayTrait;
+use Psr\Log\LoggerInterface;
 
 /**
  * Google Cloud Bigtable is Google's NoSQL Big Data database service.
@@ -144,6 +145,7 @@ public function __construct(array $config = [])
      *     @type string $appProfileId This value specifies routing for
      *           replication. **Defaults to** the "default" application profile.
      *     @type array $headers Headers to be passed with each request.
+     *     @type LoggerInterface $logger
      * }
      * @return Table
      */
diff --git a/Bigtable/src/ChunkFormatter.php b/Bigtable/src/ChunkFormatter.php
index 9a9fb1372f9f..5c7bb7b11bdd 100644
--- a/Bigtable/src/ChunkFormatter.php
+++ b/Bigtable/src/ChunkFormatter.php
@@ -140,7 +140,8 @@ function ($ex) {
                 }
                 return false;
             },
-            $this->pluck('retries', $this->options, false)
+            $this->pluck('retries', $this->options, false),
+            $this->pluck('logger', $this->options, false)
         );
     }
 
diff --git a/Bigtable/src/Table.php b/Bigtable/src/Table.php
index fe157fb9bc93..29c8e5e21f44 100644
--- a/Bigtable/src/Table.php
+++ b/Bigtable/src/Table.php
@@ -30,6 +30,7 @@
 use Google\Cloud\Bigtable\V2\RowSet;
 use Google\Cloud\Core\ArrayTrait;
 use Google\Rpc\Code;
+use Psr\Log\LoggerInterface;
 
 /**
  * A table instance can be used to read rows and to perform insert, update, and
@@ -83,6 +84,7 @@ class Table
      *           This settings only applies to {@see \Google\Cloud\Bigtable\Table::mutateRows()},
      *           {@see \Google\Cloud\Bigtable\Table::upsert()} and
      *           {@see \Google\Cloud\Bigtable\Table::readRows()}.
+     *     @type LoggerInterface $logger
      * }
      */
     public function __construct(
@@ -519,7 +521,8 @@ private function mutateRowsWithEntries(array $entries, array $options = [])
             [$this->gapicClient, 'mutateRows'],
             $argumentFunction,
             $retryFunction,
-            $this->pluck('retries', $options, false)
+            $this->pluck('retries', $options, false),
+            $this->pluck('logger', $options, false)
         );
         $message = 'partial failure';
         try {
diff --git a/Bigtable/tests/Unit/SmartRetriesTest.php b/Bigtable/tests/Unit/SmartRetriesTest.php
index 9997c7ff014d..aea0f518174c 100644
--- a/Bigtable/tests/Unit/SmartRetriesTest.php
+++ b/Bigtable/tests/Unit/SmartRetriesTest.php
@@ -37,6 +37,7 @@
 use PHPUnit\Framework\TestCase;
 use Prophecy\PhpUnit\ProphecyTrait;
 use Prophecy\Argument;
+use Psr\Log\LoggerInterface;
 
 /**
  * @group bigtable
@@ -195,7 +196,6 @@ public function testReadRowsContainsAttemptHeader()
         $attempt = 0;
         $expectedArgs = $this->options;
         $retryingApiException = $this->retryingApiException;
-        $phpunit = $this;
         $this->bigtableClient->readRows(
             self::TABLE_NAME,
             Argument::that(function ($callOptions) use (&$attempt) {
@@ -747,6 +747,63 @@ public function testMutateRowsShouldNotRetryIfAnyMutationIsNotRetryable()
         }
     }
 
+    public function testReadRowsShouldLogRetryableExeception()
+    {
+        $attempt = 0;
+        $expectedArgs = $this->options;
+        $retryingApiException = $this->retryingApiException;
+        $logger = $this->prophesize(LoggerInterface::class);
+        $logger->error('DEADLINE_EXCEEDED')
+            ->shouldBeCalledOnce();
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            Argument::that(function ($callOptions) use (&$attempt) {
+                $attemptHeader = $callOptions['headers']['bigtable-attempt'][0] ?? null;
+                ($attempt === 0)
+                    ? $this->assertNull($attemptHeader)
+                    : $this->assertSame('1', $attemptHeader);
+
+                return true;
+            })
+        )->shouldBeCalledTimes(2)
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $this->serverStream->readAll()
+            ->will(function () use (&$attempt, $retryingApiException) {
+                // throw a retriable exception on the first call
+                return 0 === $attempt++ ? throw $retryingApiException : [];
+            });
+
+        $iterator = $this->table->readRows(['logger' => $logger->reveal()]);
+        $iterator->getIterator()->current();
+    }
+
+    public function testReadRowsShouldLogNonRetryableExeception()
+    {
+        $this->expectException(ApiException::class);
+        $this->expectExceptionMessage('UNAUTHENTICATED');
+
+        $logger = $this->prophesize(LoggerInterface::class);
+        $logger->error('UNAUTHENTICATED')
+            ->shouldBeCalledOnce();
+
+        $this->bigtableClient->readRows(
+            self::TABLE_NAME,
+            $this->options
+        )->shouldBeCalledOnce()
+            ->willReturn(
+                $this->serverStream->reveal()
+            );
+
+        $this->serverStream->readAll()
+            ->willThrow($this->nonRetryingApiException);
+
+        $iterator = $this->table->readRows(['logger' => $logger->reveal()]);
+        $iterator->getIterator()->current();
+    }
+
     private function generateRowsResponse($from, $to)
     {
         $rows = [];