From d0cbb7c06ffb40e86887c859e5de9d4eb3519abd Mon Sep 17 00:00:00 2001
From: Brent Shaffer <betterbrent@google.com>
Date: Fri, 4 Oct 2024 17:50:40 -0700
Subject: [PATCH] add metric test for user refresh id tokens

---
 src/Credentials/UserRefreshCredentials.php | 10 ++++++++--
 tests/ObservabilityMetricsTest.php         | 14 ++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/Credentials/UserRefreshCredentials.php b/src/Credentials/UserRefreshCredentials.php
index d4d968c55..49ad0b7e9 100644
--- a/src/Credentials/UserRefreshCredentials.php
+++ b/src/Credentials/UserRefreshCredentials.php
@@ -57,6 +57,12 @@ class UserRefreshCredentials extends CredentialsLoader implements GetQuotaProjec
      */
     protected $quotaProject;
 
+    /**
+     * Whether this is an ID token request or an access token request. Used when
+     * building the metric header.
+     */
+    private bool $isIdTokenRequest = false;
+
     /**
      * Create a new UserRefreshCredentials.
      *
@@ -103,6 +109,7 @@ public function __construct(
         $additionalClaims = [];
         if ($targetAudience) {
             $additionalClaims = ['target_audience' => $targetAudience];
+            $this->isIdTokenRequest = true;
         }
         $this->auth = new OAuth2([
             'clientId' => $jsonKey['client_id'],
@@ -136,10 +143,9 @@ public function __construct(
      */
     public function fetchAuthToken(callable $httpHandler = null, array $metricsHeader = [])
     {
-        // We don't support id token endpoint requests as of now for User Cred
         return $this->auth->fetchAuthToken(
             $httpHandler,
-            $this->applyTokenEndpointMetrics($metricsHeader, 'at')
+            $this->applyTokenEndpointMetrics($metricsHeader, $this->isIdTokenRequest ? 'it' : 'at')
         );
     }
 
diff --git a/tests/ObservabilityMetricsTest.php b/tests/ObservabilityMetricsTest.php
index 450bfa125..002abc15e 100644
--- a/tests/ObservabilityMetricsTest.php
+++ b/tests/ObservabilityMetricsTest.php
@@ -131,10 +131,6 @@ public function testImpersonatedServiceAccountCredentials()
         $this->assertUpdateMetadata($impersonatedCred, $handler, 'imp', $handlerCalled);
     }
 
-    /**
-     * UserRefreshCredentials haven't enabled identity token support hence
-     * they don't have 'auth-request-type/it' observability metric header check.
-     */
     public function testUserRefreshCredentials()
     {
         $keyFile = __DIR__ . '/fixtures2/gcloud.json';
@@ -145,6 +141,16 @@ public function testUserRefreshCredentials()
         $this->assertUpdateMetadata($userRefreshCred, $handler, 'u', $handlerCalled);
     }
 
+    public function testUserRefreshCredentialsWithIdTokens()
+    {
+        $keyFile = __DIR__ . '/fixtures2/gcloud.json';
+        $handlerCalled = false;
+        $handler = $this->getCustomHandler('u', 'auth-request-type/it', $handlerCalled);
+
+        $userRefreshCred = new UserRefreshCredentials(null, $keyFile, 'test-target-audience');
+        $this->assertUpdateMetadata($userRefreshCred, $handler, 'u', $handlerCalled);
+    }
+
     /**
      * Invokes the 'updateMetadata' method of cred fetcher with empty metadata argument
      * and asserts for proper service api usage observability metrics header.