chromium/chrome/browser/ai/ai_rewriter_unittest.cc

// Copyright 2024 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/ai/ai_rewriter.h"

#include <memory>

#include "base/run_loop.h"
#include "base/test/bind.h"
#include "base/test/mock_callback.h"
#include "chrome/browser/ai/ai_manager_keyed_service_factory.h"
#include "chrome/browser/optimization_guide/mock_optimization_guide_keyed_service.h"
#include "chrome/browser/optimization_guide/optimization_guide_keyed_service_factory.h"
#include "chrome/test/base/chrome_render_view_host_test_harness.h"
#include "components/optimization_guide/core/mock_optimization_guide_model_executor.h"
#include "components/optimization_guide/core/optimization_guide_switches.h"
#include "components/optimization_guide/core/optimization_guide_util.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/blink/public/mojom/ai/ai_manager.mojom.h"
#include "third_party/blink/public/mojom/ai/model_streaming_responder.mojom.h"

_;
AtMost;
NiceMock;

namespace {

constexpr char kSharedContextString[] = …;
constexpr char kContextString[] = …;
constexpr char kConcatenatedContextString[] = …;
constexpr char kInputString[] = …;

class MockCreateRewriterClient
    : public blink::mojom::AIManagerCreateRewriterClient { … };

// TODO(crbug.com/358214322): Move MockResponder to a common utils file.
class MockResponder : public blink::mojom::ModelStreamingResponder { … };

class MockSupportsUserData : public base::SupportsUserData { … };

optimization_guide::OptimizationGuideModelStreamingExecutionResult
CreateExecutionResult(const std::string_view outpt, bool is_complete) { … }

optimization_guide::OptimizationGuideModelStreamingExecutionResult
CreateExecutionErrorResult(
    optimization_guide::OptimizationGuideModelExecutionError error) { … }

void CheckComposeRequestContext(
    const google::protobuf::MessageLite& request_metadata,
    const std::string& expected_context_string) { … }

void CheckComposeRequestRewriteParamsPreviousResponse(
    const google::protobuf::MessageLite& request_metadata,
    const std::string& previous_response) { … }

void CheckComposeRequestRewriteParamsTone(
    const google::protobuf::MessageLite& request_metadata,
    optimization_guide::proto::ComposeTone tone) { … }

void CheckComposeRequestRewriteParamsLength(
    const google::protobuf::MessageLite& request_metadata,
    optimization_guide::proto::ComposeLength length) { … }

void CheckComposeRequestRewriteParamsRegenerateFlag(
    const google::protobuf::MessageLite& request_metadata) { … }

}  // namespace

class AIRewriterTest : public ChromeRenderViewHostTestHarness { … };

void AIRewriterTest::RunSimpleRewriteTest(
    blink::mojom::AIRewriterTone tone,
    blink::mojom::AIRewriterLength length,
    base::OnceCallback<void(const google::protobuf::MessageLite&
                                request_metadata)> request_check_callback) { … }

void AIRewriterTest::RunRewriteOptionCombinationFailureTest(
    blink::mojom::AIRewriterTone tone,
    blink::mojom::AIRewriterLength length) { … }

TEST_F(AIRewriterTest, CreateRewriterNoService) { … }

TEST_F(AIRewriterTest, CreateRewriterModelNotAvailable) { … }

TEST_F(AIRewriterTest, CreateRewriterRetryAfterConfigNotAvailableForFeature) { … }

TEST_F(AIRewriterTest, CreateRewriterAbortAfterConfigNotAvailableForFeature) { … }

TEST_F(AIRewriterTest, ContextDestroyed) { … }

TEST_F(AIRewriterTest, RewriteRegenerate) { … }

TEST_F(AIRewriterTest, RewriteMoreCasual) { … }

TEST_F(AIRewriterTest, RewriteMoreFormal) { … }

TEST_F(AIRewriterTest, RewriteLonger) { … }

TEST_F(AIRewriterTest, RewriteShorter) { … }

TEST_F(AIRewriterTest, RewriteOptionCombinationFailureTest) { … }

TEST_F(AIRewriterTest, RewriteError) { … }

TEST_F(AIRewriterTest, RewriteMultipleResponse) { … }

TEST_F(AIRewriterTest, MultipleRewrite) { … }

TEST_F(AIRewriterTest, ResponderDisconnected) { … }

TEST_F(AIRewriterTest, RewriterDisconnected) { … }