Skip to content

Commit

Permalink
More unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@yrliou
yrliou committed Oct 30, 2024
1 parent 0db3a85 commit 763155e
Show file tree
Hide file tree
Showing 5 changed files with 228 additions and 1 deletion.
58 changes: 58 additions & 0 deletions components/ai_chat/core/browser/ai_chat_service_unittest.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,7 @@ class MockAssociatedContent
GetStagedEntriesFromContent,
(ConversationHandler::GetStagedEntriesCallback),
(override));
MOCK_METHOD(bool, HasOpenLeoPermission, (), (const, override));

base::WeakPtr<ConversationHandler::AssociatedContentDelegate> GetWeakPtr() {
return weak_ptr_factory_.GetWeakPtr();
Expand Down Expand Up @@ -492,4 +493,61 @@ TEST_P(AIChatServiceUnitTest,
run_loop.Run();
}

TEST_P(AIChatServiceUnitTest, OpenLeoWithStagedConversations_NoPermission) {
NiceMock<MockAssociatedContent> associated_content{};
ConversationHandler* conversation =
ai_chat_service_->CreateConversationHandlerForContent(
associated_content.GetContentId(), associated_content.GetWeakPtr());
auto conversation_client = CreateConversationClient(conversation);

ON_CALL(associated_content, HasOpenLeoPermission)
.WillByDefault(testing::Return(false));
EXPECT_CALL(associated_content, GetStagedEntriesFromContent).Times(0);

bool opened = false;
ai_chat_service_->OpenLeoWithStagedConversations(
associated_content.GetWeakPtr(),
base::BindLambdaForTesting([&]() { opened = true; }));
EXPECT_FALSE(opened);
testing::Mock::VerifyAndClearExpectations(&associated_content);
}

TEST_P(AIChatServiceUnitTest, OpenLeoWithStagedConversations) {
NiceMock<MockAssociatedContent> associated_content{};
ConversationHandler* conversation =
ai_chat_service_->CreateConversationHandlerForContent(
associated_content.GetContentId(), associated_content.GetWeakPtr());
auto conversation_client = CreateConversationClient(conversation);

ON_CALL(associated_content, GetStagedEntriesFromContent)
.WillByDefault(
[](ConversationHandler::GetStagedEntriesCallback callback) {
std::move(callback).Run(std::vector<SearchQuerySummary>{
SearchQuerySummary("query", "summary")});
});
ON_CALL(associated_content, HasOpenLeoPermission)
.WillByDefault(testing::Return(true));

// Allowed scheme to be associated with a conversation
ON_CALL(associated_content, GetURL())
.WillByDefault(testing::Return(GURL("https://example.com")));

// One from setting up a connected client, one from
// OpenLeoWithStagedConversations.
EXPECT_CALL(associated_content, GetStagedEntriesFromContent).Times(2);

bool opened = false;
ai_chat_service_->OpenLeoWithStagedConversations(
associated_content.GetWeakPtr(),
base::BindLambdaForTesting([&]() { opened = true; }));

base::RunLoop().RunUntilIdle();
auto& history = conversation->GetConversationHistory();
ASSERT_EQ(history.size(), 2u);
EXPECT_EQ(history[0]->text, "query");
EXPECT_EQ(history[1]->text, "summary");
EXPECT_TRUE(opened);
testing::Mock::VerifyAndClearExpectations(&associated_content);
}

} // namespace ai_chat
8 changes: 8 additions & 0 deletions components/ai_chat/core/browser/conversation_handler.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -246,6 +246,10 @@ class ConversationHandler : public mojom::ConversationHandler,
return associated_content_delegate_.get();
}

void SetRequestInProgressForTesting(bool in_progress) {
is_request_in_progress_ = in_progress;
}

protected:
// ModelService::Observer
void OnModelListUpdated() override;
Expand All @@ -263,6 +267,10 @@ class ConversationHandler : public mojom::ConversationHandler,
UpdateOrCreateLastAssistantEntry_NotDelta);
FRIEND_TEST_ALL_PREFIXES(ConversationHandlerUnitTest,
UpdateOrCreateLastAssistantEntry_NotDeltaWithSearch);
FRIEND_TEST_ALL_PREFIXES(ConversationHandlerUnitTest,
OnGetStagedEntriesFromContent);
FRIEND_TEST_ALL_PREFIXES(ConversationHandlerUnitTest,
OnGetStagedEntriesFromContent_FailedChecks);
FRIEND_TEST_ALL_PREFIXES(ConversationHandlerUnitTest, SelectedLanguage);
FRIEND_TEST_ALL_PREFIXES(PageContentRefineTest, LocalModelsUpdater);
FRIEND_TEST_ALL_PREFIXES(PageContentRefineTest, TextEmbedder);
Expand Down
128 changes: 128 additions & 0 deletions components/ai_chat/core/browser/conversation_handler_unittest.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,32 @@ class ConversationHandlerUnitTest : public testing::Test {
});
}

// Pair of text and whether it's from Brave Search SERP
void SetupHistory(std::vector<std::pair<std::string, bool>> entries) {
std::vector<mojom::ConversationTurnPtr> history;
for (size_t i = 0; i < entries.size(); i++) {
bool is_human = i % 2 == 0;

std::optional<std::vector<mojom::ConversationEntryEventPtr>> events;
if (!is_human) {
events = std::vector<mojom::ConversationEntryEventPtr>();
events->push_back(mojom::ConversationEntryEvent::NewCompletionEvent(
mojom::CompletionEvent::New(entries[i].first)));
}

auto entry = mojom::ConversationTurn::New(
is_human ? mojom::CharacterType::HUMAN
: mojom::CharacterType::ASSISTANT,
is_human ? mojom::ActionType::QUERY : mojom::ActionType::RESPONSE,
mojom::ConversationTurnVisibility::VISIBLE,
entries[i].first /* text */, std::nullopt /* selected_text */,
std::move(events), base::Time::Now(), std::nullopt /* edits */,
entries[i].second /* from_brave_search_SERP */);
history.push_back(std::move(entry));
}
conversation_handler_->SetChatHistoryForTesting(std::move(history));
}

protected:
base::test::TaskEnvironment task_environment_;
std::unique_ptr<AIChatService> ai_chat_service_;
Expand Down Expand Up @@ -1004,6 +1030,108 @@ TEST_F(ConversationHandlerUnitTest,
EXPECT_TRUE(conversation_handler_->GetConversationHistory().empty());
}

TEST_F(
ConversationHandlerUnitTest,
MaybeFetchOrClearContentStagedConversation_FetchStagedEntriesWithHistory) {
NiceMock<MockConversationHandlerClient> client(conversation_handler_.get());
ASSERT_TRUE(conversation_handler_->IsAnyClientConnected());

// MaybeFetchOrClearContentStagedConversation should clear old staged entries
// and fetch new ones.
EXPECT_CALL(*associated_content_, GetStagedEntriesFromContent).Times(1);
// One from SetupHistory and one from removing old entries and adding
// new entries in OnGetStagedEntriesFromContent.
EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(2);

// Fill history with staged and non-staged entries.
SetupHistory({{"old query" /* text */, true /*from_brave_search_SERP */},
{"old summary", "true"},
{"normal query", false},
{"normal response", false}});

// Setting mock return values for GetStagedEntriesFromContent.
SetAssociatedContentStagedEntries(/*empty=*/false, /*multi=*/true);

conversation_handler_->MaybeFetchOrClearContentStagedConversation();
task_environment_.RunUntilIdle();

testing::Mock::VerifyAndClearExpectations(associated_content_.get());
testing::Mock::VerifyAndClearExpectations(&client);

auto& history = conversation_handler_->GetConversationHistory();
ASSERT_EQ(history.size(), 6u);
EXPECT_FALSE(history[0]->from_brave_search_SERP);
EXPECT_EQ(history[0]->text, "normal query");
EXPECT_FALSE(history[1]->from_brave_search_SERP);
EXPECT_EQ(history[1]->text, "normal response");
EXPECT_TRUE(history[2]->from_brave_search_SERP);
EXPECT_EQ(history[2]->text, "query");
EXPECT_TRUE(history[3]->from_brave_search_SERP);
EXPECT_EQ(history[3]->text, "summary");
EXPECT_TRUE(history[4]->from_brave_search_SERP);
EXPECT_EQ(history[4]->text, "query2");
EXPECT_TRUE(history[5]->from_brave_search_SERP);
EXPECT_EQ(history[5]->text, "summary2");
}

TEST_F(ConversationHandlerUnitTest,
OnGetStagedEntriesFromContent_FailedChecks) {
// No staged entries would be added if a request is in progress.
conversation_handler_->SetRequestInProgressForTesting(true);
std::vector<SearchQuerySummary> entries = {{"query", "summary"},
{"query2", "summary2"}};
conversation_handler_->OnGetStagedEntriesFromContent(entries);
task_environment_.RunUntilIdle();
EXPECT_EQ(conversation_handler_->GetConversationHistory().size(), 0u);

// No staged entries if should_send_page_contents_ is false.
conversation_handler_->SetRequestInProgressForTesting(false);
conversation_handler_->SetShouldSendPageContents(false);
conversation_handler_->OnGetStagedEntriesFromContent(entries);
task_environment_.RunUntilIdle();
EXPECT_EQ(conversation_handler_->GetConversationHistory().size(), 0u);

// No staged entries if user opt-out.
conversation_handler_->SetShouldSendPageContents(true);
EmulateUserOptedOut();
conversation_handler_->OnGetStagedEntriesFromContent(entries);
task_environment_.RunUntilIdle();
EXPECT_EQ(conversation_handler_->GetConversationHistory().size(), 0u);
}

TEST_F(ConversationHandlerUnitTest, OnGetStagedEntriesFromContent) {
NiceMock<MockConversationHandlerClient> client(conversation_handler_.get());
ASSERT_TRUE(conversation_handler_->IsAnyClientConnected());

EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(2);
// Fill history with staged and non-staged entries.
SetupHistory({{"q1" /* text */, true /*from_brave_search_SERP */},
{"s1", "true"},
{"q2", false},
{"r1", false}});

std::vector<SearchQuerySummary> entries = {{"query", "summary"},
{"query2", "summary2"}};
conversation_handler_->OnGetStagedEntriesFromContent(entries);
task_environment_.RunUntilIdle();
testing::Mock::VerifyAndClearExpectations(&client);

auto& history = conversation_handler_->GetConversationHistory();
ASSERT_EQ(history.size(), 6u);
EXPECT_FALSE(history[0]->from_brave_search_SERP);
EXPECT_EQ(history[0]->text, "q2");
EXPECT_FALSE(history[1]->from_brave_search_SERP);
EXPECT_EQ(history[1]->text, "r1");
EXPECT_TRUE(history[2]->from_brave_search_SERP);
EXPECT_EQ(history[2]->text, "query");
EXPECT_TRUE(history[3]->from_brave_search_SERP);
EXPECT_EQ(history[3]->text, "summary");
EXPECT_TRUE(history[4]->from_brave_search_SERP);
EXPECT_EQ(history[4]->text, "query2");
EXPECT_TRUE(history[5]->from_brave_search_SERP);
EXPECT_EQ(history[5]->text, "summary2");
}

TEST_F(ConversationHandlerUnitTest_OptedOut,
MaybeFetchOrClearSearchQuerySummary_NotOptedIn) {
// Content will have staged entries, but we want to make sure that
Expand Down
5 changes: 4 additions & 1 deletion components/ai_chat/core/common/BUILD.gn
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,10 @@ component("common") {
if (!is_ios) {
source_set("unit_tests") {
testonly = true
sources = [ "pref_names_unittest.cc" ]
sources = [
"pref_names_unittest.cc",
"utils_unittest.cc",
]

deps = [
"//base/test:test_support",
Expand Down
30 changes: 30 additions & 0 deletions components/ai_chat/core/common/utils_unittest.cc
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
/* Copyright (c) 2024 The Brave Authors. All rights reserved.
* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this file,
* You can obtain one at https://mozilla.org/MPL/2.0/. */

#include "brave/components/ai_chat/core/common/utils.h"

#include "testing/gtest/include/gtest/gtest.h"
#include "url/gurl.h"

namespace ai_chat {

TEST(AIChatCommonUtilsUnitTest, IsBraveSearchURL) {
EXPECT_TRUE(IsBraveSearchURL(GURL("https://search.brave.com")));
EXPECT_FALSE(IsBraveSearchURL(GURL("http://search.brave.com")));
EXPECT_FALSE(IsBraveSearchURL(GURL("https://test.brave.com/")));
EXPECT_FALSE(IsBraveSearchURL(GURL("https://brave.com/")));
EXPECT_FALSE(IsBraveSearchURL(GURL()));
}

TEST(AIChatCommonUtilsUnitTest, IsOpenLeoButtonFromBraveSearchURL) {
EXPECT_TRUE(IsOpenLeoButtonFromBraveSearchURL(GURL("https://search.brave.com/leo#5566")));
EXPECT_FALSE(IsOpenLeoButtonFromBraveSearchURL(GURL()));
EXPECT_FALSE(IsOpenLeoButtonFromBraveSearchURL(GURL("https://search.brave.com")));
EXPECT_FALSE(IsOpenLeoButtonFromBraveSearchURL(GURL("https://search.brave.com/leo")));
EXPECT_FALSE(IsOpenLeoButtonFromBraveSearchURL(GURL("https://search.brave.com/leo#")));
EXPECT_FALSE(IsOpenLeoButtonFromBraveSearchURL(GURL("https://brave.com/leo#5566")));
}

}

0 comments on commit 763155e

Please sign in to comment.