// 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 "chromeos/ash/components/data_migration/data_migration.h"
#include <cstdint>
#include <optional>
#include <utility>
#include "base/base_paths.h"
#include "base/containers/flat_map.h"
#include "base/containers/flat_set.h"
#include "base/files/file_util.h"
#include "base/files/scoped_temp_dir.h"
#include "base/functional/bind.h"
#include "base/memory/raw_ptr.h"
#include "base/ranges/algorithm.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/stringprintf.h"
#include "base/test/bind.h"
#include "base/test/run_until.h"
#include "base/test/scoped_path_override.h"
#include "base/test/task_environment.h"
#include "chromeos/ash/components/data_migration/constants.h"
#include "chromeos/ash/components/data_migration/testing/connection_barrier.h"
#include "chromeos/ash/components/data_migration/testing/fake_nearby_connections.h"
#include "chromeos/ash/components/data_migration/testing/fake_nearby_process_manager.h"
#include "chromeos/ash/components/nearby/common/connections_manager/nearby_connections_manager.h"
#include "chromeos/ash/components/nearby/common/connections_manager/nearby_connections_manager_impl.h"
#include "chromeos/ash/services/nearby/public/mojom/nearby_connections_types.mojom.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
namespace data_migration {
namespace {
constexpr char kRemoteEndpointId[] = "test-remote-endpoint";
// TODO(esum): Move common test harness logic to a dedicated class. Some of this
// is shared with other unit tests.
class DataMigrationTest : public ::testing::Test {
public:
DataMigrationTest()
: nearby_process_manager_(kRemoteEndpointId),
data_migration_(std::make_unique<NearbyConnectionsManagerImpl>(
&nearby_process_manager_,
kServiceId)) {}
void SetUp() override {
ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
ASSERT_TRUE(base::CreateDirectory(GetFilePayloadDirectory()));
home_dir_override_.emplace(base::DIR_HOME, temp_dir_.GetPath());
nearby_process_manager_.fake_nearby_connections()
.set_local_to_remote_payload_listener(base::BindRepeating(
&DataMigrationTest::RespondToCtsMessage, base::Unretained(this)));
}
void TearDown() override {
static_cast<KeyedService*>(&data_migration_)->Shutdown();
}
base::FilePath GetFilePayloadDirectory() {
return temp_dir_.GetPath().Append(kPayloadTargetDir);
}
base::FilePath BuildFilePayloadName(int64_t payload_id) {
return base::FilePath(base::StringPrintf("payload_%ld", payload_id));
}
base::FilePath BuildFilePayloadPath(int64_t payload_id) {
return GetFilePayloadDirectory().Append(BuildFilePayloadName(payload_id));
}
void RespondToCtsMessage(
::nearby::connections::mojom::PayloadPtr cts_payload) {
ASSERT_TRUE(cts_payload);
ASSERT_TRUE(cts_payload->content->is_bytes());
const std::vector<uint8_t>& cts_bytes =
cts_payload->content->get_bytes()->bytes;
int64_t file_payload_id_to_transmit = 0;
ASSERT_TRUE(
base::StringToInt64(std::string(cts_bytes.begin(), cts_bytes.end()),
&file_payload_id_to_transmit));
ASSERT_TRUE(
requested_file_payload_ids_.contains(file_payload_id_to_transmit));
ASSERT_TRUE(nearby_process_manager_.fake_nearby_connections().SendFile(
file_payload_id_to_transmit,
&expected_file_content_[file_payload_id_to_transmit]));
}
void SendRts(int64_t file_payload_id) {
static int64_t g_rts_payload_id_assigner = 1000;
ASSERT_TRUE(
nearby_process_manager_.fake_nearby_connections().SendBytesPayload(
/*payload_id=*/g_rts_payload_id_assigner,
base::NumberToString(file_payload_id)));
++g_rts_payload_id_assigner;
requested_file_payload_ids_.insert(file_payload_id);
}
bool FileIsReady(int64_t payload_id) {
base::FilePath file_path = BuildFilePayloadPath(payload_id);
int64_t file_size_in_bytes = 0;
return base::PathExists(file_path) &&
base::GetFileSize(file_path, &file_size_in_bytes) &&
file_size_in_bytes >=
nearby_process_manager_.fake_nearby_connections()
.test_file_size_in_bytes();
}
base::flat_set</*payload_id*/ int64_t> requested_file_payload_ids_;
base::flat_map</*payload_id*/ int64_t, std::vector<uint8_t>>
expected_file_content_;
base::test::TaskEnvironment task_environment_;
base::ScopedTempDir temp_dir_;
std::optional<base::ScopedPathOverride> home_dir_override_;
FakeNearbyProcessManager nearby_process_manager_;
DataMigration data_migration_;
};
TEST_F(DataMigrationTest, CompletesAllFileTransfers) {
nearby_process_manager_.fake_nearby_connections()
.set_connection_established_listener(base::BindLambdaForTesting([this]() {
SendRts(/*file_payload_id=*/1);
SendRts(/*file_payload_id=*/2);
SendRts(/*file_payload_id=*/3);
}));
data_migration_.StartAdvertising();
ASSERT_TRUE(base::test::RunUntil([this]() {
// All expected files have been written to disc.
return base::ranges::all_of(
requested_file_payload_ids_,
[this](int64_t payload_id) { return FileIsReady(payload_id); });
}));
// For extra safety, flush any pending tasks to ensure `DataMigration` is
// completely idle before checking the results.
task_environment_.RunUntilIdle();
EXPECT_EQ(base::ReadFileToBytes(BuildFilePayloadPath(/*payload_id=*/1)),
expected_file_content_.at(/*payload_id=*/1));
EXPECT_EQ(base::ReadFileToBytes(BuildFilePayloadPath(/*payload_id=*/2)),
expected_file_content_.at(/*payload_id=*/2));
EXPECT_EQ(base::ReadFileToBytes(BuildFilePayloadPath(/*payload_id=*/3)),
expected_file_content_.at(/*payload_id=*/3));
}
TEST_F(DataMigrationTest, HandlesDisconnect) {
// The transfer should be in progress when the remote device is disconnected.
nearby_process_manager_.fake_nearby_connections().SetFinalFilePayloadStatus(
FakeNearbyConnections::PayloadStatus::kInProgress, /*payload_id=*/1);
nearby_process_manager_.fake_nearby_connections()
.set_connection_established_listener(base::BindLambdaForTesting(
[this]() { SendRts(/*file_payload_id=*/1); }));
data_migration_.StartAdvertising();
// Run until the transfers starts, then disconnect.
ASSERT_TRUE(base::test::RunUntil([this]() {
return base::PathExists(BuildFilePayloadPath(/*payload_id=*/1));
}));
// The transfer should succeed the second time.
nearby_process_manager_.fake_nearby_connections().SetFinalFilePayloadStatus(
FakeNearbyConnections::PayloadStatus::kSuccess, /*payload_id=*/1);
ASSERT_TRUE(nearby_process_manager_.fake_nearby_connections()
.SimulateRemoteDisconnect());
// Advertising/discovery should be retried and succeed again. This time,
// the file transfer should succeed.
ASSERT_TRUE(
base::test::RunUntil([this]() { return FileIsReady(/*payload_id=*/1); }));
// For extra safety, flush any pending tasks to ensure `DataMigration` is
// completely idle before checking the results.
task_environment_.RunUntilIdle();
EXPECT_EQ(base::ReadFileToBytes(BuildFilePayloadPath(/*payload_id=*/1)),
expected_file_content_.at(/*payload_id=*/1));
}
// Verifies that `DataMigration` can be destroyed gracefully mid-transfer. Does
// not have any real expectation other than the code doesn't crash.
TEST_F(DataMigrationTest, ShutsDownMidTransfer) {
// The transfer should be in progress when shutting down.
nearby_process_manager_.fake_nearby_connections().SetFinalFilePayloadStatus(
FakeNearbyConnections::PayloadStatus::kInProgress, /*payload_id=*/1);
nearby_process_manager_.fake_nearby_connections()
.set_connection_established_listener(base::BindLambdaForTesting(
[this]() { SendRts(/*file_payload_id=*/1); }));
data_migration_.StartAdvertising();
// Run until the transfers starts, then proceed to the test
// harness's `TearDown()`.
ASSERT_TRUE(base::test::RunUntil([this]() {
return base::PathExists(BuildFilePayloadPath(/*payload_id=*/1));
}));
}
} // namespace
} // namespace data_migration