chromeos/ash/components/data_migration/device_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 "chromeos/ash/components/data_migration/device.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/location.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/task/task_traits.h"
#include "base/task/thread_pool.h"
#include "base/test/bind.h"
#include "base/test/run_until.h"
#include "base/test/scoped_path_override.h"
#include "base/test/scoped_run_loop_timeout.h"
#include "base/test/task_environment.h"
#include "base/timer/timer.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";

class DeviceTest : public ::testing::Test {
 public:
  DeviceTest()
      : nearby_process_manager_(kRemoteEndpointId),
        nearby_connections_manager_(&nearby_process_manager_, kServiceId) {}

  void SetUp() override {
    ASSERT_TRUE(temp_dir_.CreateUniqueTempDir());
    home_dir_override_.emplace(base::DIR_HOME, temp_dir_.GetPath());
    ConnectionBarrier connection_barrier(&nearby_connections_manager_);
    NearbyConnection* nearby_connection = connection_barrier.Wait();
    ASSERT_TRUE(nearby_connection);
    ASSERT_TRUE(base::CreateDirectory(GetFilePayloadDirectory()));
    nearby_process_manager_.fake_nearby_connections()
        .set_local_to_remote_payload_listener(base::BindRepeating(
            &DeviceTest::RespondToCtsMessage, base::Unretained(this)));
    device_.emplace(nearby_connection, &nearby_connections_manager_);
  }

  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) {
    // Must not intersect with file payload ids, which are assigned starting at
    // "1" in test cases below.
    static int64_t g_rts_payload_id_assigner = 1000000;
    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_;
  NearbyConnectionsManagerImpl nearby_connections_manager_;
  std::optional<Device> device_;
};

TEST_F(DeviceTest, TransmitsAllFiles) {
  // Remote device send RTS for some files.
  SendRts(/*file_payload_id=*/1);
  SendRts(/*file_payload_id=*/2);
  SendRts(/*file_payload_id=*/3);

  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 the `Device` 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));
}

// Verifies that the Device can be destroyed gracefully mid-transfer. Does not
// have any real expectation other than the code doesn't crash.
TEST_F(DeviceTest, ShutsDownMidTransfer) {
  constexpr int64_t kTotalNumFilesToTransmit = 5;
  // Remote device send RTS for some files.
  for (int64_t file_payload_id = 1; file_payload_id <= kTotalNumFilesToTransmit;
       ++file_payload_id) {
    nearby_process_manager_.fake_nearby_connections().SetFinalFilePayloadStatus(
        ::nearby::connections::mojom::PayloadStatus::kInProgress,
        file_payload_id);
    SendRts(file_payload_id);
  }

  ASSERT_TRUE(base::test::RunUntil([this]() {
    // Any file has been written to disc (exit any time mid-transfer).
    for (auto payload_id : requested_file_payload_ids_) {
      if (base::PathExists(BuildFilePayloadPath(payload_id))) {
        return true;
      }
    }
    return false;
  }));

  ASSERT_LT(base::ComputeDirectorySize(temp_dir_.GetPath()),
            kTotalNumFilesToTransmit *
                nearby_process_manager_.fake_nearby_connections()
                    .test_file_size_in_bytes());
  device_.reset();
}

TEST_F(DeviceTest, ContinuesTransmittingFilesAfterFailure) {
  nearby_process_manager_.fake_nearby_connections().SetFinalFilePayloadStatus(
      ::nearby::connections::mojom::PayloadStatus::kFailure, /*payload_id=*/2);

  // Remote device send RTS for some files.
  SendRts(/*file_payload_id=*/1);
  SendRts(/*file_payload_id=*/2);
  SendRts(/*file_payload_id=*/3);

  ASSERT_TRUE(base::test::RunUntil([this]() {
    return FileIsReady(/*payload_id=*/1) && FileIsReady(/*payload_id=*/3);
  }));
  // For extra safety, flush any pending tasks to ensure the `Device` 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=*/3)),
            expected_file_content_.at(/*payload_id=*/3));
}

TEST_F(DeviceTest, FileStressTest) {
  constexpr int64_t kTotalNumFilesToTransmit = 5000;
  // Reduce these numbers from their defaults otherwise the test takes 3-4x
  // longer.
  constexpr int64_t kFileSizeInBytes = 1;
  constexpr int64_t kNumChunksPerFile = 1;

  // Empirically, this test takes about 15 seconds.
  base::test::ScopedRunLoopTimeout test_timeout(FROM_HERE, base::Minutes(1));

  nearby_process_manager_.fake_nearby_connections().set_test_file_size_in_bytes(
      kFileSizeInBytes);
  nearby_process_manager_.fake_nearby_connections().set_test_file_num_chunks(
      kNumChunksPerFile);

  // Remote device send RTS for some files.
  for (int64_t payload_id = 1; payload_id <= kTotalNumFilesToTransmit;
       ++payload_id) {
    SendRts(payload_id);
  }

  base::RepeatingClosure test_complete_signal = task_environment_.QuitClosure();
  // The `test_completion_check` is expensive, so it's run on a blocking thread,
  // freeing the main thread for the test itself.
  base::RepeatingClosure test_completion_check =
      base::BindLambdaForTesting([this, test_complete_signal]() {
        if (base::ComputeDirectorySize(GetFilePayloadDirectory()) >=
            kTotalNumFilesToTransmit * kFileSizeInBytes) {
          test_complete_signal.Run();
        }
      });
  scoped_refptr<base::SequencedTaskRunner> test_completion_sequence =
      base::ThreadPool::CreateSequencedTaskRunner({base::MayBlock()});
  auto test_completion_checker = std::make_unique<base::RepeatingTimer>();
  test_completion_sequence->PostTask(
      FROM_HERE, base::BindLambdaForTesting([&test_completion_checker,
                                             test_completion_check]() {
        constexpr base::TimeDelta kTestCompletionCheckInterval =
            base::Seconds(1);
        test_completion_checker->Start(FROM_HERE, kTestCompletionCheckInterval,
                                       test_completion_check);
      }));

  task_environment_.RunUntilQuit();
  test_completion_sequence->DeleteSoon(FROM_HERE,
                                       std::move(test_completion_checker));
  // Flush any pending ThreadPool tasks scheduled internally by `DataMigration`
  // before exiting.
  task_environment_.RunUntilIdle();
}

}  // namespace
}  // namespace data_migration