chrome/browser/devtools/aida_client_unittest.cc - chromium/src - Git at Google

 // Copyright 2023 The Chromium Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 #include "chrome/browser/devtools/aida_client.h"

 #include <memory>
 #include <utility>

 #include "base/functional/bind.h"
 #include "base/test/metrics/histogram_tester.h"
 #include "chrome/browser/signin/identity_test_environment_profile_adaptor.h"
 #include "chrome/test/base/testing_profile.h"
 #include "content/public/browser/network_service_instance.h"
 #include "content/public/test/browser_task_environment.h"
 #include "content/public/test/test_utils.h"
 #include "net/test/embedded_test_server/embedded_test_server.h"
 #include "services/network/network_service.h"
 #include "services/network/public/mojom/network_context_client.mojom.h"
 #include "services/network/test/test_shared_url_loader_factory.h"
 #include "testing/gtest/include/gtest/gtest.h"

 namespace {
 using net::test_server::BasicHttpResponse;
 using net::test_server::HttpRequest;
 using net::test_server::HttpResponse;
 }  // namespace

 const char kEmail[] = "alice@example.com";
 const char kEndpointPath[] = "/foo";
 const char kScope[] = "bar";
 const char kRequest[] =
     R"({"input": "What does this code do: 1+1", "client": "GENERAL"})";
 const char kResponse[] =
     R"([{"textChunk":{"text":"The function `foo()` takes no arguments and returns nothing."}}])";

 class AidaClientTest : public testing::Test {
  public:
   AidaClientTest()
       : task_environment_(content::BrowserTaskEnvironment::IO_MAINLOOP),
         profile_(IdentityTestEnvironmentProfileAdaptor::
                      CreateProfileForIdentityTestEnvironment()),
         identity_test_env_adaptor_(
             std::make_unique<IdentityTestEnvironmentProfileAdaptor>(
                 profile_.get())),
         identity_test_env_(identity_test_env_adaptor_->identity_test_env()) {
     content::GetNetworkService();
     content::RunAllPendingInMessageLoop(content::BrowserThread::IO);

     shared_url_loader_factory_ =
         base::MakeRefCounted<network::TestSharedURLLoaderFactory>(
             network::NetworkService::GetNetworkServiceForTesting());

     identity_test_env_->MakePrimaryAccountAvailable(
         kEmail, signin::ConsentLevel::kSync);
   }

  protected:
   content::BrowserTaskEnvironment task_environment_;
   std::unique_ptr<network::mojom::NetworkContextClient> network_context_client_;
   scoped_refptr<network::TestSharedURLLoaderFactory> shared_url_loader_factory_;
   net::EmbeddedTestServer test_server_;
   std::unique_ptr<TestingProfile> profile_;
   std::unique_ptr<IdentityTestEnvironmentProfileAdaptor>
       identity_test_env_adaptor_;
   signin::IdentityTestEnvironment* identity_test_env_;
   base::HistogramTester histogram_tester_;
 };

 class Delegate {
  public:
   Delegate() = default;

   std::unique_ptr<HttpResponse> HandleRequest(const HttpRequest& request) {
     request_ = request.content;
     authorization_header_ =
         request.headers.at(net::HttpRequestHeaders::kAuthorization);

     auto http_response = std::make_unique<BasicHttpResponse>();
     http_response->set_code(api_response_code_);
     http_response->set_content(api_response_);
     http_response->set_content_type("application/json");
     return http_response;
   }

   void FinishCallback(base::RunLoop* run_loop, const std::string& response) {
     response_ = response;
     if (run_loop) {
       run_loop->Quit();
     }
   }

   std::string request_;
   std::string api_response_ = kResponse;
   net::HttpStatusCode api_response_code_ = net::HTTP_OK;
   std::string response_;
   std::string authorization_header_;
 };

 constexpr char kOAuthToken[] = "5678";

 TEST_F(AidaClientTest, DoesNothingIfNoScope) {
   Delegate delegate;
   test_server_.RegisterRequestHandler(base::BindRepeating(
       &Delegate::HandleRequest, base::Unretained(&delegate)));

   AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
   aida_client.OverrideAidaEndpointAndScopeForTesting("", "");
   aida_client.DoConversation(
       kRequest, base::BindOnce(&Delegate::FinishCallback,
                                base::Unretained(&delegate), nullptr));
   EXPECT_EQ("", delegate.request_);
   EXPECT_EQ(R"([{"error": "AIDA scope is not configured"}])",
             delegate.response_);
 }

 TEST_F(AidaClientTest, FailsIfNotAuthorized) {
   base::RunLoop run_loop;
   Delegate delegate;

   AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
   aida_client.OverrideAidaEndpointAndScopeForTesting("https://example.com/foo",
                                                      kScope);
   aida_client.DoConversation(
       kRequest, base::BindOnce(&Delegate::FinishCallback,
                                base::Unretained(&delegate), &run_loop));
   identity_test_env_->WaitForAccessTokenRequestIfNecessaryAndRespondWithError(
       GoogleServiceAuthError(GoogleServiceAuthError::REQUEST_CANCELED));

   EXPECT_EQ("", delegate.request_);
   EXPECT_EQ(
       R"([{"error": "Cannot get OAuth credentials", "detail": "Request canceled."}])",
       delegate.response_);
 }

 TEST_F(AidaClientTest, Succeeds) {
   base::RunLoop run_loop;
   Delegate delegate;
   test_server_.RegisterRequestHandler(base::BindRepeating(
       &Delegate::HandleRequest, base::Unretained(&delegate)));

   ASSERT_TRUE(test_server_.Start());

   GURL endpoint_url = test_server_.GetURL(kEndpointPath);
   AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
   aida_client.OverrideAidaEndpointAndScopeForTesting(endpoint_url.spec(),
                                                      kScope);
   aida_client.DoConversation(
       kRequest, base::BindOnce(&Delegate::FinishCallback,
                                base::Unretained(&delegate), &run_loop));
   identity_test_env_
       ->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
           kOAuthToken, base::Time::Now() + base::Seconds(10),
           std::string() /*id_token*/, signin::ScopeSet{kScope});
   run_loop.Run();

   EXPECT_EQ(kRequest, delegate.request_);
   EXPECT_EQ(kResponse, delegate.response_);
   histogram_tester_.ExpectTotalCount("DevTools.AidaResponseTime", 1);
 }

 TEST_F(AidaClientTest, ReusesOAuthToken) {
   base::RunLoop run_loop;
   Delegate delegate;
   test_server_.RegisterRequestHandler(base::BindRepeating(
       &Delegate::HandleRequest, base::Unretained(&delegate)));

   ASSERT_TRUE(test_server_.Start());

   GURL endpoint_url = test_server_.GetURL(kEndpointPath);
   AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
   aida_client.OverrideAidaEndpointAndScopeForTesting(endpoint_url.spec(),
                                                      kScope);
   aida_client.DoConversation(
       kRequest, base::BindOnce(&Delegate::FinishCallback,
                                base::Unretained(&delegate), &run_loop));
   identity_test_env_
       ->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
           kOAuthToken, base::Time::Now() + base::Seconds(10),
           std::string() /*id_token*/, signin::ScopeSet{kScope});
   run_loop.Run();

   EXPECT_EQ(kRequest, delegate.request_);
   EXPECT_EQ(kResponse, delegate.response_);
   std::string authorization_header = delegate.authorization_header_;

   const char kAnotherRequest[] = "another request";
   const char kAnotherResponse[] = "another response";
   delegate.api_response_ = kAnotherResponse;
   base::RunLoop run_loop2;
   aida_client.DoConversation(
       kAnotherRequest, base::BindOnce(&Delegate::FinishCallback,
                                       base::Unretained(&delegate), &run_loop2));
   run_loop2.Run();
   EXPECT_EQ(kAnotherRequest, delegate.request_);
   EXPECT_EQ(kAnotherResponse, delegate.response_);
   EXPECT_EQ(authorization_header, delegate.authorization_header_);
 }

 TEST_F(AidaClientTest, RefetchesTokenIfUnauthorized) {
   base::RunLoop run_loop;
   Delegate delegate;
   test_server_.RegisterRequestHandler(base::BindRepeating(
       &Delegate::HandleRequest, base::Unretained(&delegate)));

   ASSERT_TRUE(test_server_.Start());

   GURL endpoint_url = test_server_.GetURL(kEndpointPath);
   AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
   aida_client.OverrideAidaEndpointAndScopeForTesting(endpoint_url.spec(),
                                                      kScope);
   aida_client.DoConversation(
       kRequest, base::BindOnce(&Delegate::FinishCallback,
                                base::Unretained(&delegate), &run_loop));
   identity_test_env_
       ->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
           kOAuthToken, base::Time::Now() + base::Seconds(10),
           std::string() /*id_token*/, signin::ScopeSet{kScope});
   run_loop.Run();

   EXPECT_EQ(kRequest, delegate.request_);
   EXPECT_EQ(kResponse, delegate.response_);
   std::string authorization_header = delegate.authorization_header_;

   delegate.api_response_code_ = net::HTTP_UNAUTHORIZED;
   base::RunLoop run_loop2;
   const char kAnotherRequest[] = "another request";
   const char kAnotherResponse[] = "another response";
   const char kAnotherOAuthToken[] = "another token";

   aida_client.DoConversation(
       kAnotherRequest, base::BindOnce(&Delegate::FinishCallback,
                                       base::Unretained(&delegate), &run_loop2));
   identity_test_env_
       ->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
           kAnotherOAuthToken, base::Time::Now() + base::Seconds(10),
           std::string() /*id_token*/, signin::ScopeSet{kScope});
   delegate.api_response_code_ = net::HTTP_OK;
   delegate.api_response_ = kAnotherResponse;

   run_loop2.Run();
   EXPECT_EQ(kAnotherRequest, delegate.request_);
   EXPECT_EQ(kAnotherResponse, delegate.response_);
   EXPECT_NE(authorization_header, delegate.authorization_header_);
 }
	// Copyright 2023 The Chromium Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	#include "chrome/browser/devtools/aida_client.h"

	#include <memory>
	#include <utility>

	#include "base/functional/bind.h"
	#include "base/test/metrics/histogram_tester.h"
	#include "chrome/browser/signin/identity_test_environment_profile_adaptor.h"
	#include "chrome/test/base/testing_profile.h"
	#include "content/public/browser/network_service_instance.h"
	#include "content/public/test/browser_task_environment.h"
	#include "content/public/test/test_utils.h"
	#include "net/test/embedded_test_server/embedded_test_server.h"
	#include "services/network/network_service.h"
	#include "services/network/public/mojom/network_context_client.mojom.h"
	#include "services/network/test/test_shared_url_loader_factory.h"
	#include "testing/gtest/include/gtest/gtest.h"

	namespace {
	using net::test_server::BasicHttpResponse;
	using net::test_server::HttpRequest;
	using net::test_server::HttpResponse;
	} // namespace

	const char kEmail[] = "alice@example.com";
	const char kEndpointPath[] = "/foo";
	const char kScope[] = "bar";
	const char kRequest[] =
	R"({"input": "What does this code do: 1+1", "client": "GENERAL"})";
	const char kResponse[] =
	R"([{"textChunk":{"text":"The function `foo()` takes no arguments and returns nothing."}}])";

	class AidaClientTest : public testing::Test {
	public:
	AidaClientTest()
	: task_environment_(content::BrowserTaskEnvironment::IO_MAINLOOP),
	profile_(IdentityTestEnvironmentProfileAdaptor::
	CreateProfileForIdentityTestEnvironment()),
	identity_test_env_adaptor_(
	std::make_unique<IdentityTestEnvironmentProfileAdaptor>(
	profile_.get())),
	identity_test_env_(identity_test_env_adaptor_->identity_test_env()) {
	content::GetNetworkService();
	content::RunAllPendingInMessageLoop(content::BrowserThread::IO);

	shared_url_loader_factory_ =
	base::MakeRefCounted<network::TestSharedURLLoaderFactory>(
	network::NetworkService::GetNetworkServiceForTesting());

	identity_test_env_->MakePrimaryAccountAvailable(
	kEmail, signin::ConsentLevel::kSync);
	}

	protected:
	content::BrowserTaskEnvironment task_environment_;
	std::unique_ptr<network::mojom::NetworkContextClient> network_context_client_;
	scoped_refptr<network::TestSharedURLLoaderFactory> shared_url_loader_factory_;
	net::EmbeddedTestServer test_server_;
	std::unique_ptr<TestingProfile> profile_;
	std::unique_ptr<IdentityTestEnvironmentProfileAdaptor>
	identity_test_env_adaptor_;
	signin::IdentityTestEnvironment* identity_test_env_;
	base::HistogramTester histogram_tester_;
	};

	class Delegate {
	public:
	Delegate() = default;

	std::unique_ptr<HttpResponse> HandleRequest(const HttpRequest& request) {
	request_ = request.content;
	authorization_header_ =
	request.headers.at(net::HttpRequestHeaders::kAuthorization);

	auto http_response = std::make_unique<BasicHttpResponse>();
	http_response->set_code(api_response_code_);
	http_response->set_content(api_response_);
	http_response->set_content_type("application/json");
	return http_response;
	}

	void FinishCallback(base::RunLoop* run_loop, const std::string& response) {
	response_ = response;
	if (run_loop) {
	run_loop->Quit();
	}
	}

	std::string request_;
	std::string api_response_ = kResponse;
	net::HttpStatusCode api_response_code_ = net::HTTP_OK;
	std::string response_;
	std::string authorization_header_;
	};

	constexpr char kOAuthToken[] = "5678";

	TEST_F(AidaClientTest, DoesNothingIfNoScope) {
	Delegate delegate;
	test_server_.RegisterRequestHandler(base::BindRepeating(
	&Delegate::HandleRequest, base::Unretained(&delegate)));

	AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
	aida_client.OverrideAidaEndpointAndScopeForTesting("", "");
	aida_client.DoConversation(
	kRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), nullptr));
	EXPECT_EQ("", delegate.request_);
	EXPECT_EQ(R"([{"error": "AIDA scope is not configured"}])",
	delegate.response_);
	}

	TEST_F(AidaClientTest, FailsIfNotAuthorized) {
	base::RunLoop run_loop;
	Delegate delegate;

	AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
	aida_client.OverrideAidaEndpointAndScopeForTesting("https://example.com/foo",
	kScope);
	aida_client.DoConversation(
	kRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), &run_loop));
	identity_test_env_->WaitForAccessTokenRequestIfNecessaryAndRespondWithError(
	GoogleServiceAuthError(GoogleServiceAuthError::REQUEST_CANCELED));

	EXPECT_EQ("", delegate.request_);
	EXPECT_EQ(
	R"([{"error": "Cannot get OAuth credentials", "detail": "Request canceled."}])",
	delegate.response_);
	}

	TEST_F(AidaClientTest, Succeeds) {
	base::RunLoop run_loop;
	Delegate delegate;
	test_server_.RegisterRequestHandler(base::BindRepeating(
	&Delegate::HandleRequest, base::Unretained(&delegate)));

	ASSERT_TRUE(test_server_.Start());

	GURL endpoint_url = test_server_.GetURL(kEndpointPath);
	AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
	aida_client.OverrideAidaEndpointAndScopeForTesting(endpoint_url.spec(),
	kScope);
	aida_client.DoConversation(
	kRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), &run_loop));
	identity_test_env_
	->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
	kOAuthToken, base::Time::Now() + base::Seconds(10),
	std::string() /id_token/, signin::ScopeSet{kScope});
	run_loop.Run();

	EXPECT_EQ(kRequest, delegate.request_);
	EXPECT_EQ(kResponse, delegate.response_);
	histogram_tester_.ExpectTotalCount("DevTools.AidaResponseTime", 1);
	}

	TEST_F(AidaClientTest, ReusesOAuthToken) {
	base::RunLoop run_loop;
	Delegate delegate;
	test_server_.RegisterRequestHandler(base::BindRepeating(
	&Delegate::HandleRequest, base::Unretained(&delegate)));

	ASSERT_TRUE(test_server_.Start());

	GURL endpoint_url = test_server_.GetURL(kEndpointPath);
	AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
	aida_client.OverrideAidaEndpointAndScopeForTesting(endpoint_url.spec(),
	kScope);
	aida_client.DoConversation(
	kRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), &run_loop));
	identity_test_env_
	->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
	kOAuthToken, base::Time::Now() + base::Seconds(10),
	std::string() /id_token/, signin::ScopeSet{kScope});
	run_loop.Run();

	EXPECT_EQ(kRequest, delegate.request_);
	EXPECT_EQ(kResponse, delegate.response_);
	std::string authorization_header = delegate.authorization_header_;

	const char kAnotherRequest[] = "another request";
	const char kAnotherResponse[] = "another response";
	delegate.api_response_ = kAnotherResponse;
	base::RunLoop run_loop2;
	aida_client.DoConversation(
	kAnotherRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), &run_loop2));
	run_loop2.Run();
	EXPECT_EQ(kAnotherRequest, delegate.request_);
	EXPECT_EQ(kAnotherResponse, delegate.response_);
	EXPECT_EQ(authorization_header, delegate.authorization_header_);
	}

	TEST_F(AidaClientTest, RefetchesTokenIfUnauthorized) {
	base::RunLoop run_loop;
	Delegate delegate;
	test_server_.RegisterRequestHandler(base::BindRepeating(
	&Delegate::HandleRequest, base::Unretained(&delegate)));

	ASSERT_TRUE(test_server_.Start());

	GURL endpoint_url = test_server_.GetURL(kEndpointPath);
	AidaClient aida_client(profile_.get(), shared_url_loader_factory_);
	aida_client.OverrideAidaEndpointAndScopeForTesting(endpoint_url.spec(),
	kScope);
	aida_client.DoConversation(
	kRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), &run_loop));
	identity_test_env_
	->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
	kOAuthToken, base::Time::Now() + base::Seconds(10),
	std::string() /id_token/, signin::ScopeSet{kScope});
	run_loop.Run();

	EXPECT_EQ(kRequest, delegate.request_);
	EXPECT_EQ(kResponse, delegate.response_);
	std::string authorization_header = delegate.authorization_header_;

	delegate.api_response_code_ = net::HTTP_UNAUTHORIZED;
	base::RunLoop run_loop2;
	const char kAnotherRequest[] = "another request";
	const char kAnotherResponse[] = "another response";
	const char kAnotherOAuthToken[] = "another token";

	aida_client.DoConversation(
	kAnotherRequest, base::BindOnce(&Delegate::FinishCallback,
	base::Unretained(&delegate), &run_loop2));
	identity_test_env_
	->WaitForAccessTokenRequestIfNecessaryAndRespondWithTokenForScopes(
	kAnotherOAuthToken, base::Time::Now() + base::Seconds(10),
	std::string() /id_token/, signin::ScopeSet{kScope});
	delegate.api_response_code_ = net::HTTP_OK;
	delegate.api_response_ = kAnotherResponse;

	run_loop2.Run();
	EXPECT_EQ(kAnotherRequest, delegate.request_);
	EXPECT_EQ(kAnotherResponse, delegate.response_);
	EXPECT_NE(authorization_header, delegate.authorization_header_);
	}