chromium/chrome/browser/ui/webui/ash/lobster/lobster_page_handler_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/ui/webui/ash/lobster/lobster_page_handler.h"

#include <string_view>
#include <vector>

#include "ash/public/cpp/lobster/lobster_result.h"
#include "ash/public/cpp/lobster/lobster_session.h"
#include "base/base64.h"
#include "base/files/file_path.h"
#include "base/files/scoped_temp_dir.h"
#include "base/test/test_future.h"
#include "chrome/browser/download/chrome_download_manager_delegate.h"
#include "chrome/browser/download/download_core_service_factory.h"
#include "chrome/browser/download/download_core_service_impl.h"
#include "chrome/browser/download/download_prefs.h"
#include "chrome/test/base/testing_profile.h"
#include "content/public/test/browser_task_environment.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace ash {

namespace {

constexpr std::string_view kRawBytes1 = "a1b2c3";
constexpr std::string_view kRawBytes2 = "d4e5f6";

class FakeLobsterSession : public LobsterSession {
 public:
  FakeLobsterSession(const LobsterResult& result,
                     bool commit_or_download_status,
                     bool feedback_submission_status)
      : result_(result),
        commit_or_download_status_(commit_or_download_status),
        feedback_submission_status_(feedback_submission_status) {}
  ~FakeLobsterSession() override = default;

  void DownloadCandidate(int candidate_id,
                         const base::FilePath& file_path,
                         StatusCallback callback) override {
    std::move(callback).Run(commit_or_download_status_);
  }
  void CommitAsInsert(int candidate_id, StatusCallback callback) override {
    std::move(callback).Run(commit_or_download_status_);
  }
  void CommitAsDownload(int candidate_id,
                        const base::FilePath& file_path,
                        StatusCallback callback) override {
    std::move(callback).Run(commit_or_download_status_);
  }
  void RequestCandidates(const std::string& query,
                         int num_candidates,
                         RequestCandidatesCallback callback) override {
    std::move(callback).Run(result_);
  }
  void PreviewFeedback(int candidate_id,
                       LobsterPreviewFeedbackCallback callback) override {}
  bool SubmitFeedback(int candidate_id,
                      const std::string& description) override {
    return feedback_submission_status_;
  }

 private:
  LobsterResult result_;
  bool commit_or_download_status_;
  bool feedback_submission_status_;
};

class LobsterPageHandlerTest : public testing::Test {
 public:
  void SetUp() override {
    DownloadCoreServiceFactory::GetForBrowserContext(&profile_)
        ->SetDownloadManagerDelegateForTesting(
            std::make_unique<ChromeDownloadManagerDelegate>(&profile_));

    // Use a temporary directory for downloads.
    ASSERT_TRUE(download_dir_.CreateUniqueTempDir());
    DownloadPrefs* prefs =
        DownloadPrefs::FromDownloadManager(profile_.GetDownloadManager());
    prefs->SetDownloadPath(download_dir_.GetPath());
    prefs->SkipSanitizeDownloadTargetPathForTesting();
  }
  TestingProfile& profile() { return profile_; }

 protected:
  content::BrowserTaskEnvironment task_environment_;
  base::ScopedTempDir download_dir_;

 private:
  TestingProfile profile_;
};

TEST_F(LobsterPageHandlerTest,
       RequestCandidatesReturnsImagesInCorrectJpegFormat) {
  std::vector<LobsterImageCandidate> image_candidates = {
      LobsterImageCandidate(/*id=*/0, /*image_bytes=*/kRawBytes1.data(),
                            /*seed=*/20,
                            /*query=*/"a nice strawberry"),
      LobsterImageCandidate(/*id=*/1, /*image_bytes=*/kRawBytes2.data(),
                            /*seed=*/21,
                            /*query=*/"a nice strawberry")};
  FakeLobsterSession session(std::move(image_candidates),
                             /*commit_or_download_status=*/true,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<lobster::mojom::ResponsePtr> future;

  page_handler.RequestCandidates("a nice strawberry", 2, future.GetCallback());

  EXPECT_TRUE(future.Get()->is_candidates());

  auto& actual_candidates = future.Get()->get_candidates();

  EXPECT_EQ(actual_candidates.size(), 2u);
  EXPECT_EQ(actual_candidates[0]->id, 0u);
  EXPECT_EQ(actual_candidates[0]->data_url,
            GURL(base::StrCat(
                {"data:image/jpeg;base64,", base::Base64Encode(kRawBytes1)})));
  EXPECT_EQ(actual_candidates[1]->id, 1u);
  EXPECT_EQ(actual_candidates[1]->data_url,
            GURL(base::StrCat(
                {"data:image/jpeg;base64,", base::Base64Encode(kRawBytes2)})));
}

TEST_F(LobsterPageHandlerTest, RequestCandidatesReturnsError) {
  FakeLobsterSession session(
      base::unexpected(
          LobsterError(LobsterErrorCode::kInvalidArgument, "dummy error")),
      /*commit_or_download_status=*/false, /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<lobster::mojom::ResponsePtr> future;

  page_handler.RequestCandidates("a nice strawberry", 2, future.GetCallback());

  EXPECT_TRUE(future.Get()->is_error());

  auto& actual_error = future.Get()->get_error();

  EXPECT_EQ(actual_error->code, LobsterErrorCode::kInvalidArgument);
  EXPECT_EQ(actual_error->message, "dummy error");
}

TEST_F(LobsterPageHandlerTest, DownloadCandidateSucceeds) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/true,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.DownloadCandidate(1, future.GetCallback());

  EXPECT_TRUE(future.Get());
}

TEST_F(LobsterPageHandlerTest, DownloadCandidateFails) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/false,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.DownloadCandidate(/*id=*/1, future.GetCallback());

  EXPECT_FALSE(future.Get());
}

TEST_F(LobsterPageHandlerTest, CommitAsDownloadSucceeds) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/true,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.CommitAsDownload(/*id=*/1, future.GetCallback());

  EXPECT_TRUE(future.Get());
}

TEST_F(LobsterPageHandlerTest, CommitAsDownloadFails) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/false,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.CommitAsDownload(/*id=*/1, future.GetCallback());

  EXPECT_FALSE(future.Get());
}

TEST_F(LobsterPageHandlerTest, CommitAsInsertSucceeds) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/true,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.CommitAsInsert(/*id=*/1, future.GetCallback());

  EXPECT_TRUE(future.Get());
}

TEST_F(LobsterPageHandlerTest, CommitAsInsertFails) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/false,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.CommitAsInsert(/*id=*/1, future.GetCallback());

  EXPECT_FALSE(future.Get());
}

TEST_F(LobsterPageHandlerTest, SubmitFeedbackFails) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/false,
                             /*feedback_submission_status=*/false);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.SubmitFeedback(/*id=*/1, /*description=*/"dummy description",
                              future.GetCallback());

  EXPECT_FALSE(future.Get());
}

TEST_F(LobsterPageHandlerTest, SubmitFeedbackSucceeds) {
  FakeLobsterSession session({}, /*commit_or_download_status=*/false,
                             /*feedback_submission_status=*/true);
  LobsterPageHandler page_handler = LobsterPageHandler(&session, &profile());
  base::test::TestFuture<bool> future;

  page_handler.SubmitFeedback(/*id=*/1, /*description=*/"dummy description",
                              future.GetCallback());

  EXPECT_TRUE(future.Get());
}

}  // namespace
}  // namespace ash