/*
 * Copyright 2019 Bloomberg Finance LP
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_merklize.h>
#include <buildboxcommon_permissions.h>
#include <buildboxcommon_temporarydirectory.h>
#include <buildboxcommon_temporaryfile.h>
#include <buildboxcommon_timeutils.h>
#include <buildboxcommontest_utils.h>
#include <gtest/gtest.h>
#include <ThreadPool.h>

#include <build/bazel/remote/execution/v2/remote_execution_mock.grpc.pb.h>
#include <build/buildgrid/local_cas_mock.grpc.pb.h>
#include <google/bytestream/bytestream.pb.h>
#include <google/bytestream/bytestream_mock.grpc.pb.h>
#include <grpcpp/test/mock_stream.h>

#include <algorithm>
#include <fstream>
#include <memory>
#include <sys/stat.h>

using namespace buildboxcommon;
using namespace testing;

namespace {
const auto digestFunctionInitializer = []() {
    buildboxcommon::DigestGenerator::init();
    return 0;
}();
} // namespace

const int64_t MAX_BATCH_SIZE_BYTES = 2UL * 1024;

class StubsFixture : public ::testing::Test {
    /**
     * Fixture that provides mock grpc stubs for instantiating the CAS
     * clients in this file's tests.
     */
  protected:
    std::shared_ptr<google::bytestream::MockByteStreamStub> bytestreamClient;
    std::shared_ptr<MockContentAddressableStorageStub> casClient;
    std::shared_ptr<MockLocalContentAddressableStorageStub> localCasClient;
    std::shared_ptr<MockCapabilitiesStub> capabilitiesClient;

    StubsFixture()
    {
        bytestreamClient =
            std::make_shared<google::bytestream::MockByteStreamStub>();
        casClient = std::make_shared<MockContentAddressableStorageStub>();
        localCasClient =
            std::make_shared<MockLocalContentAddressableStorageStub>();
        capabilitiesClient = std::make_shared<MockCapabilitiesStub>();
    }
};

TEST_F(StubsFixture, InitTest)
{
    auto grpcClient = std::make_shared<GrpcClient>();
    CASClient client(grpcClient);
    auto cacheCapabilities = std::make_unique<CacheCapabilities>();
    cacheCapabilities->set_max_batch_total_size_bytes(64);
    cacheCapabilities->add_digest_functions(
        DigestFunction_Value::DigestFunction_Value_SHA256);
    ServerCapabilities serverCapabilities;
    serverCapabilities.set_allocated_cache_capabilities(
        cacheCapabilities.release());

    EXPECT_CALL(*capabilitiesClient, GetCapabilities(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(serverCapabilities),
                        Return(grpc::Status::OK)));
    client.init(bytestreamClient, casClient, localCasClient,
                capabilitiesClient);

    ASSERT_TRUE(grpcClient->instanceName().empty());
}

TEST_F(StubsFixture, InitCapabilitiesDidntReturnOk)
{
    /**
     * This test ensures that we're still ok if the capabilities endpoint
     * doesn't support the capabilities request or fails in some other way
     * (we just use a default value).
     */
    auto grpcClient = std::make_shared<GrpcClient>();
    CASClient client(grpcClient);
    EXPECT_CALL(*capabilitiesClient, GetCapabilities(_, _, _))
        .WillOnce(Return(
            grpc::Status(grpc::UNIMPLEMENTED, "method not found for test")));
    client.init(bytestreamClient, casClient, localCasClient,
                capabilitiesClient);
}

TEST_F(StubsFixture, InitUnsupportedDigestFunction)
{
    /* Client requests SHA-256, server only supports SHA-512 */
    auto grpcClient = std::make_shared<GrpcClient>();
    CASClient client(grpcClient);
    auto cacheCapabilities = std::make_unique<CacheCapabilities>();
    cacheCapabilities->set_max_batch_total_size_bytes(64);
    cacheCapabilities->add_digest_functions(
        DigestFunction_Value::DigestFunction_Value_SHA512);
    ServerCapabilities serverCapabilities;
    serverCapabilities.set_allocated_cache_capabilities(
        cacheCapabilities.release());

    EXPECT_CALL(*capabilitiesClient, GetCapabilities(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(serverCapabilities),
                        Return(grpc::Status::OK)));
    EXPECT_THROW(client.init(bytestreamClient, casClient, localCasClient,
                             capabilitiesClient),
                 std::runtime_error);
}

using MockBytestreamReader =
    grpc::testing::MockClientReader<google::bytestream::ReadResponse>;
using MockBytestreamWriter =
    grpc::testing::MockClientWriter<google::bytestream::WriteRequest>;
using MockGetTreeReaderWriter = grpc::testing::MockClientReaderWriter<
    typename build::buildgrid::StageTreeRequest,
    typename build::buildgrid::StageTreeResponse>;
using MockGetTreeReader = grpc::testing::MockClientReader<
    typename build::bazel::remote::execution::v2::GetTreeResponse>;

class ClientTestFixture : public StubsFixture, public CASClient {
    /**
     * Fixture that provides a pre-instantiated client, as well as several
     * objects to be passed as arguments and returned from mocks.
     *
     * Inherits from the fixture that provides stubs.
     */
  protected:
    const std::string client_instance_name = "CasTestInstance123";
    const std::string content = "password";
    Digest digest;
    TemporaryFile tmpfile;
    google::bytestream::ReadResponse readResponse;
    google::bytestream::WriteResponse writeResponse;

    ClientTestFixture(int64_t max_batch_size_bytes = MAX_BATCH_SIZE_BYTES)
        : CASClient(std::make_shared<GrpcClient>())
    {
        this->init(bytestreamClient, casClient, localCasClient, nullptr,
                   max_batch_size_bytes);
        this->getGrpcClient()->setInstanceName(client_instance_name);
        this->getGrpcClient()->setRetryLimit(1);
    }
};

TEST_F(ClientTestFixture, FetchStringTest)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest = DigestGenerator::hash(content);

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));

    ReadRequest request;
    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(reader.release())));

    // Setting a new instance name with the client's setter:
    const std::string instance_name = "newInstanceName!";
    this->getGrpcClient()->setInstanceName(instance_name);
    ASSERT_EQ(this->getGrpcClient()->instanceName(), instance_name);

    GrpcClient::RequestStats requestStats;
    EXPECT_EQ(this->fetchString(digest, &requestStats), readResponse.data());
    EXPECT_EQ(request.resource_name().substr(0, instance_name.size()),
              instance_name);
    ASSERT_EQ(requestStats.d_grpcRetryCount, 0);
}

TEST_F(ClientTestFixture, FetchStringEmptyResponse)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data("");
    digest = DigestGenerator::hash("");

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;
    EXPECT_EQ(this->fetchString(digest), readResponse.data());
}

TEST_F(ClientTestFixture, FetchStringSizeMismatch)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest.set_size_bytes(99999);

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish())
        .WillOnce(
            Return(grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, "")));
    auto _ = reader.release();
    (void)_;

    EXPECT_THROW(this->fetchString(digest), std::runtime_error);
}

TEST_F(ClientTestFixture, FetchStringServerError)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest.set_size_bytes(static_cast<int64_t>(content.length()));

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::StatusCode::NOT_FOUND, "Digest not found!")));
    auto _ = reader.release();
    (void)_;

    GrpcClient::RequestStats requestStats;
    EXPECT_THROW(this->fetchString(digest), std::runtime_error);
    ASSERT_EQ(requestStats.d_grpcRetryCount, 0);
}

TEST_F(ClientTestFixture, FetchStringServerRetryableError)
{
    readResponse.set_data(content);
    digest.set_size_bytes(static_cast<int64_t>(content.length()));

    ASSERT_EQ(getGrpcClient()->retryLimit(), 1);

    auto reader1 = std::make_unique<MockBytestreamReader>();
    auto reader2 = std::make_unique<MockBytestreamReader>();

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader1.get()))
        .WillOnce(Return(reader2.get()));

    EXPECT_CALL(*reader1, Read(_)).WillOnce(Return(false));
    EXPECT_CALL(*reader1, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                      "Something is wrong right now.")));
    auto _1 = reader1.release();
    (void)_1;

    EXPECT_CALL(*reader2, Read(_)).WillOnce(Return(false));
    EXPECT_CALL(*reader2, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                      "Something is still wrong.")));
    auto _2 = reader2.release();
    (void)_2;

    GrpcClient::RequestStats requestStats;
    ASSERT_THROW(this->fetchString(digest, &requestStats), GrpcError);
    ASSERT_EQ(requestStats.d_grpcRetryCount, getGrpcClient()->retryLimit());
}

TEST_F(ClientTestFixture, FindMissingBlobsSuccessful)
{
    Digest missingDigest;
    missingDigest.set_hash("missing-hash");

    FindMissingBlobsResponse response;
    response.add_missing_blob_digests()->CopyFrom(missingDigest);

    Digest presentDigest;
    presentDigest.set_hash("present-hash");

    const std::vector<Digest> findMissingList = {missingDigest, presentDigest};

    EXPECT_CALL(*casClient.get(), FindMissingBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(response), Return(grpc::Status::OK)));

    GrpcClient::RequestStats requestStats;
    const auto missingBlobsResponse =
        findMissingBlobs(findMissingList, &requestStats);
    EXPECT_EQ(requestStats.d_grpcRetryCount, 0);
    ASSERT_EQ(missingBlobsResponse.size(), 1);
    ASSERT_EQ(missingBlobsResponse.front(), missingDigest);
}

TEST_F(ClientTestFixture, FindMissingBlobsRetryableError)
{
    Digest missingDigest;
    missingDigest.set_hash("missing-hash");

    FindMissingBlobsResponse response;
    response.add_missing_blob_digests()->CopyFrom(missingDigest);

    Digest presentDigest;
    presentDigest.set_hash("present-hash");

    const std::vector<Digest> findMissingList = {missingDigest, presentDigest};

    const auto errorStatus = grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                          "Something is wrong right now.");
    EXPECT_CALL(*casClient.get(), FindMissingBlobs(_, _, _))
        .Times(getGrpcClient()->retryLimit() + 1)
        .WillRepeatedly(Return(errorStatus));

    GrpcClient::RequestStats requestStats;
    try {
        findMissingBlobs(findMissingList, &requestStats);
    }
    catch (const GrpcError &e) {
        EXPECT_EQ(e.status.error_code(), errorStatus.error_code());
        EXPECT_EQ(e.status.error_message(), errorStatus.error_message());
    }
    catch (...) {
        FAIL() << "findMissingBlobs() threw Unexpected exception type.\n";
    }

    ASSERT_EQ(requestStats.d_grpcRetryCount, getGrpcClient()->retryLimit());
}

TEST_F(ClientTestFixture, DownloadTest)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest = DigestGenerator::hash(content);

    ReadRequest request;
    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(reader.get())));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    this->download(tmpfile.fd(), digest);

    tmpfile.close();

    std::ifstream in(tmpfile.name());
    std::stringstream buffer;
    buffer << in.rdbuf();

    EXPECT_EQ(buffer.str(), readResponse.data());
    EXPECT_EQ(request.resource_name().substr(0, client_instance_name.size()),
              client_instance_name);
}

TEST_F(ClientTestFixture, DownloadFdNotWritable)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest.set_size_bytes(static_cast<int64_t>(content.length()));

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)));
    auto _ = reader.release();
    (void)_;

    EXPECT_THROW(this->download(-1, digest), std::system_error);
}

TEST_F(ClientTestFixture, DownloadSizeMismatch)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest.set_size_bytes(99999999);

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    EXPECT_THROW(this->download(tmpfile.fd(), digest), std::runtime_error);
}

TEST_F(ClientTestFixture, DownloadHashMismatch)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest.set_hash("invalid-hash");
    digest.set_size_bytes(static_cast<int64_t>(content.length()));

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    EXPECT_THROW(this->download(tmpfile.fd(), digest), std::runtime_error);
}

TEST_F(ClientTestFixture, DownloadServerError)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    readResponse.set_data(content);
    digest.set_size_bytes(static_cast<int64_t>(content.length()));

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::NOT_FOUND,
                                      "Do not have the digest (fatal)")));
    auto _ = reader.release();
    (void)_;

    EXPECT_THROW(this->download(tmpfile.fd(), digest), GrpcError);
}

TEST_F(ClientTestFixture, DownloadRetryableServerError)
{
    readResponse.set_data(content);
    digest.set_size_bytes(static_cast<int64_t>(content.length()));

    ASSERT_EQ(getGrpcClient()->retryLimit(), 1);

    auto reader1 = std::make_unique<MockBytestreamReader>();
    auto reader2 = std::make_unique<MockBytestreamReader>();

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader1.get()))
        .WillOnce(Return(reader2.get()));

    EXPECT_CALL(*reader1, Read(_)).WillOnce(Return(false));
    EXPECT_CALL(*reader1, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                      "Something is wrong right now.")));
    auto _1 = reader1.release();
    (void)_1;

    EXPECT_CALL(*reader2, Read(_)).WillOnce(Return(false));
    EXPECT_CALL(*reader2, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                      "Something is still wrong.")));
    auto _2 = reader2.release();
    (void)_2;

    GrpcClient::RequestStats requestStats;
    EXPECT_THROW(this->download(tmpfile.fd(), digest, &requestStats),
                 GrpcError);
    EXPECT_EQ(requestStats.d_grpcRetryCount, getGrpcClient()->retryLimit());
}

TEST_F(ClientTestFixture, UploadStringTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    WriteRequest request;
    EXPECT_CALL(*writer, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(content, digest);
    EXPECT_EQ(request.resource_name().substr(0, client_instance_name.size()),
              client_instance_name);
}

TEST_F(ClientTestFixture, UploadStringCommittedSizeMismatch)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes() - 1);
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    WriteRequest request;
    EXPECT_CALL(*writer, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    ASSERT_THROW(this->upload(content, digest), std::runtime_error);
}

TEST_F(ClientTestFixture, UploadLargeStringTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 3 * 1024 * 1024;
    std::string content = std::string(contentLength, 'f');
    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).Times(3).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(content, digest);
}

TEST_F(ClientTestFixture, UploadExactStringTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 1024 * 1024;
    std::string content = std::string(contentLength, 'f');
    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(content, digest);
}

TEST_F(ClientTestFixture, UploadJustLargerThanExactStringTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 1024 * 1024 + 1;
    std::string content = std::string(contentLength, 'f');
    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).Times(2).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(content, digest);
}

TEST_F(ClientTestFixture, UploadJustSmallerThanExactStringTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 1024 * 1024 - 1;
    std::string content = std::string(contentLength, 'f');
    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(content, digest);
}

TEST_F(ClientTestFixture, UploadStringSizeMismatch)
{
    digest.set_size_bytes(9999999999999999);
    digest.set_hash("fakehash");

    EXPECT_THROW(this->upload(content, digest), std::logic_error);
}

TEST_F(ClientTestFixture, UploadAlreadyExistingString)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const std::string data("This blob is already present in the remote.");
    const Digest digest = DigestGenerator::hash(data);

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(false));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish())
        .WillOnce(Return(grpc::Status(grpc::Status::OK)));
    auto _ = writer.release();
    (void)_;

    ASSERT_NO_THROW(this->upload(data, digest));
}

TEST_F(ClientTestFixture, UploadStringDidntReturnOk)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.get()));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::FAILED_PRECONDITION, "failing for test")));
    auto _ = writer.release();
    (void)_;

    EXPECT_THROW(this->upload(content, digest), GrpcError);
}

TEST_F(ClientTestFixture, UploadStringRetriableError)
{
    auto writer1 = std::make_unique<MockBytestreamWriter>();
    auto writer2 = std::make_unique<MockBytestreamWriter>();

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer1.get()))
        .WillOnce(Return(writer2.get()));

    EXPECT_CALL(*writer1, Write(_, _)).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer1, WritesDone()).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer1, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::UNAVAILABLE, "Something is wrong right now")));
    auto _1 = writer1.release();
    (void)_1;

    EXPECT_CALL(*writer2, Write(_, _)).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer2, WritesDone()).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer2, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::UNAVAILABLE, "Something is wrong right now")));
    auto _2 = writer2.release();
    (void)_2;

    EXPECT_THROW(this->upload(content, DigestGenerator::hash(content)),
                 GrpcError);
}

TEST_F(ClientTestFixture, FileTooLargeToBatchUpload)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const auto data = std::string(3 * MAX_BATCH_SIZE_BYTES, '_');
    const Digest digest = DigestGenerator::hash(data);

    CASClient::UploadRequest request(digest, data);
    std::vector<CASClient::UploadRequest> requests = {request};

    // Expecting it to fall back to a bytestream Write():
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.get()));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    this->uploadBlobs(requests);
    auto _ = writer.release();
    (void)_;
}

TEST_F(ClientTestFixture, UploadBlobs)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const std::vector<std::string> payload = {
        "a", "b", std::string(2 * MAX_BATCH_SIZE_BYTES, 'x'), "c"};
    const std::vector<std::string> hashes = {"hash0", "hash1", "hash2",
                                             "hash3"};

    // Creating list of requests...
    std::vector<CASClient::UploadRequest> requests;
    for (unsigned i = 0; i < payload.size(); i++) {
        Digest digest;
        digest.set_hash(hashes[i]);
        digest.set_size_bytes(static_cast<int64_t>(payload[i].size()));

        CASClient::UploadRequest request(digest, payload[i]);

        requests.push_back(request);
    }
    ASSERT_EQ(requests.size(), payload.size());

    // We expect digests {0, 1, 3} to be uploaded with BatchUpdateBlobs().
    BatchUpdateBlobsResponse response;
    for (unsigned i = 0; i < payload.size(); i++) {
        if (i == 2)
            continue;
        auto entry = response.add_responses();
        entry->mutable_digest()->CopyFrom(requests[i].digest);
    }

    BatchUpdateBlobsRequest request;
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));
    // And digest in index 2 with the Bytestream API:
    writeResponse.set_committed_size(requests[2].digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    const auto failed_uploads = this->uploadBlobs(requests);
    ASSERT_TRUE(failed_uploads.empty());

    // The client sends the correct instance name:
    EXPECT_EQ(request.instance_name(), client_instance_name);
}

TEST_F(ClientTestFixture, UploadBlobsThreaded)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const std::vector<std::string> payload = {
        "a", "b", std::string(2 * MAX_BATCH_SIZE_BYTES, 'x'), "c"};
    const std::vector<std::string> hashes = {"hash0", "hash1", "hash2",
                                             "hash3"};

    // Creating list of requests...
    std::vector<CASClient::UploadRequest> requests;
    for (unsigned i = 0; i < payload.size(); i++) {
        Digest digest;
        digest.set_hash(hashes[i]);
        digest.set_size_bytes(static_cast<int64_t>(payload[i].size()));

        CASClient::UploadRequest request(digest, payload[i]);

        requests.push_back(request);
    }
    ASSERT_EQ(requests.size(), payload.size());

    // We expect digests {0, 1, 3} to be uploaded with BatchUpdateBlobs().
    BatchUpdateBlobsResponse response;
    for (unsigned i = 0; i < payload.size(); i++) {
        if (i == 2)
            continue;
        auto entry = response.add_responses();
        entry->mutable_digest()->CopyFrom(requests[i].digest);
    }

    BatchUpdateBlobsRequest request;
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));
    // And digest in index 2 with the Bytestream API:
    writeResponse.set_committed_size(requests[2].digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    ThreadPool pool(2, "");
    const auto failed_uploads = this->uploadBlobs(requests, &pool);
    ASSERT_TRUE(failed_uploads.empty());

    // The client sends the correct instance name:
    EXPECT_EQ(request.instance_name(), client_instance_name);
}

TEST_F(ClientTestFixture, UploadBlobsReturnsFailures)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const std::vector<std::string> payload = {
        "a", std::string(2 * MAX_BATCH_SIZE_BYTES, 'x')};
    const std::vector<std::string> hashes = {"hash0", "hash1"};

    // Creating list of requests...
    std::vector<CASClient::UploadRequest> requests;
    for (unsigned i = 0; i < payload.size(); i++) {
        Digest digest;
        digest.set_hash(hashes[i]);
        digest.set_size_bytes(static_cast<int64_t>(payload[i].size()));

        CASClient::UploadRequest request(digest, payload[i]);

        requests.push_back(request);
    }
    ASSERT_EQ(requests.size(), payload.size());

    // Both requests will fail with:
    const auto errorStatus = grpc::Status(grpc::StatusCode::INTERNAL,
                                          "Could not write data in CAS.");

    BatchUpdateBlobsResponse response;
    auto entry = response.add_responses();
    entry->mutable_digest()->CopyFrom(requests[0].digest);

    entry->mutable_status()->set_code(errorStatus.error_code());
    entry->mutable_status()->set_message(errorStatus.error_message());
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(response), Return(grpc::Status::OK)));

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.get()));
    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(errorStatus));
    auto _ = writer.release();
    (void)_;

    const auto failed_uploads = this->uploadBlobs(requests);

    ASSERT_EQ(failed_uploads.size(), 2);
    ASSERT_EQ(std::count(hashes.cbegin(), hashes.cend(),
                         failed_uploads[0].digest.hash()),
              1);
    ASSERT_EQ(std::count(hashes.cbegin(), hashes.cend(),
                         failed_uploads[1].digest.hash()),
              1);
}

TEST_F(ClientTestFixture, UploadBlobsReturnsFatalErrors)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const std::vector<std::string> payload = {
        "a", std::string(2 * MAX_BATCH_SIZE_BYTES, 'x')};
    const std::vector<std::string> hashes = {"hash0", "hash1"};

    // Creating list of requests...
    std::vector<CASClient::UploadRequest> requests;
    for (unsigned i = 0; i < payload.size(); i++) {
        Digest digest;
        digest.set_hash(hashes[i]);
        digest.set_size_bytes(static_cast<int64_t>(payload[i].size()));

        CASClient::UploadRequest request(digest, payload[i]);

        requests.push_back(request);
    }
    ASSERT_EQ(requests.size(), payload.size());

    // Both requests will fail with:
    const auto errorStatus =
        grpc::Status(grpc::StatusCode::PERMISSION_DENIED, "Permission denied");

    BatchUpdateBlobsResponse response;
    auto entry = response.add_responses();
    entry->mutable_digest()->CopyFrom(requests[0].digest);

    entry->mutable_status()->set_code(errorStatus.error_code());
    entry->mutable_status()->set_message(errorStatus.error_message());
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(response), Return(errorStatus)));

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.get()));
    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(errorStatus));
    auto _ = writer.release();
    (void)_;

    const auto failed_uploads = this->uploadBlobs(requests);

    ASSERT_EQ(failed_uploads.size(), 2);
    ASSERT_EQ(std::count(hashes.cbegin(), hashes.cend(),
                         failed_uploads[0].digest.hash()),
              1);
    ASSERT_EQ(std::count(hashes.cbegin(), hashes.cend(),
                         failed_uploads[1].digest.hash()),
              1);
    ASSERT_EQ(failed_uploads.at(0).status.error_code(),
              grpc::StatusCode::PERMISSION_DENIED);
    ASSERT_EQ(failed_uploads.at(1).status.error_code(),
              grpc::StatusCode::PERMISSION_DENIED);
}

TEST_F(ClientTestFixture, UploadBlobsEmpty)
{
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _)).Times(0);

    this->uploadBlobs({});
}

TEST_F(ClientTestFixture, CaptureDirectory)
{
    const std::string path_to_capture = "/path/to/stage";
    std::vector<std::string> paths = {path_to_capture};
    const std::string property = "mtime";
    const std::vector<std::string> properties = {property};

    CaptureTreeResponse response;
    auto entry = response.add_responses();
    entry->set_path(path_to_capture);
    entry->mutable_tree_digest()->CopyFrom(DigestGenerator::hash("tree-blob"));
    entry->mutable_status()->set_code(grpc::StatusCode::OK);

    CaptureTreeRequest request;

    EXPECT_CALL(*localCasClient.get(), CaptureTree(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));

    const CaptureTreeResponse returned_response =
        this->captureTree(paths, properties, false, 0);

    // Checking that the request has the data we expect:
    ASSERT_EQ(request.path_size(), 1);
    ASSERT_EQ(request.path(0), path_to_capture);
    ASSERT_EQ(request.node_properties_size(), 1);
    ASSERT_EQ(request.node_properties(0), property);
    ASSERT_FALSE(request.bypass_local_cache());
    ASSERT_EQ(request.instance_name(), this->getGrpcClient()->instanceName());

    // Checking the response:
    ASSERT_EQ(response.responses_size(), 1);
    ASSERT_EQ(response.responses(0).path(), path_to_capture);
    ASSERT_EQ(response.responses(0).status().code(), grpc::StatusCode::OK);
}

TEST_F(ClientTestFixture, CaptureDirectoryErrorThrows)
{
    std::vector<std::string> paths = {std::string("/path/to/stage")};
    std::vector<std::string> properties = {std::string("mtime")};

    CaptureTreeResponse response;
    auto entry = response.add_responses();
    entry->set_path("/dev/null");
    entry->mutable_tree_digest()->CopyFrom(DigestGenerator::hash("tree-blob"));
    entry->mutable_status()->set_code(grpc::StatusCode::OK);

    CaptureTreeRequest request;

    // The retry logic throws after running out of tries:
    EXPECT_CALL(*localCasClient.get(), CaptureTree(_, _, _))
        .WillRepeatedly(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                              Return(grpc::Status(grpc::StatusCode::UNKNOWN,
                                                  "Something went wrong."))));

    ASSERT_THROW(this->captureTree(paths, properties, false, 0),
                 std::runtime_error);
}

TEST_F(ClientTestFixture, CaptureFiles)
{

    const std::vector<std::string> files_to_capture = {
        "/path/to/stage/file1.txt", "/path/to/stage/file2.txt"};

    const std::string property = "mtime";
    const auto mtime = TimeUtils::now();

    // Response that the server will return to the client:
    CaptureFilesResponse response;
    auto entry1 = response.add_responses();
    entry1->set_path(files_to_capture[0]);
    entry1->mutable_digest()->CopyFrom(
        DigestGenerator::hash("file1.txt-contents"));
    entry1->mutable_status()->set_code(grpc::StatusCode::OK);
    auto properties1 = entry1->mutable_node_properties();
    properties1->mutable_mtime()->CopyFrom(mtime);

    auto entry2 = response.add_responses();
    entry2->set_path(files_to_capture[1]);
    entry2->mutable_digest()->CopyFrom(
        DigestGenerator::hash("file2.txt-contents"));
    entry2->mutable_status()->set_code(grpc::StatusCode::OK);
    auto properties2 = entry2->mutable_node_properties();
    properties2->mutable_mtime()->CopyFrom(mtime);

    CaptureFilesRequest request;

    EXPECT_CALL(*localCasClient.get(), CaptureFiles(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));

    const std::vector<std::string> properties = {property};

    const auto returned_responses =
        this->captureFiles(files_to_capture, properties, false);

    // Checking that the request issued contains the data we expect:
    const std::set<std::string> files_to_capture_set(files_to_capture.cbegin(),
                                                     files_to_capture.cend());

    ASSERT_EQ(request.path_size(), 2);
    ASSERT_EQ(files_to_capture_set.count(request.path(0)), 1);
    ASSERT_EQ(files_to_capture_set.count(request.path(1)), 1);
    ASSERT_EQ(request.node_properties_size(), 1);
    ASSERT_EQ(request.node_properties(0), property);

    ASSERT_FALSE(request.bypass_local_cache());

    ASSERT_EQ(request.instance_name(), this->getGrpcClient()->instanceName());

    // Checking that the response returned by the client; it should match the
    // one issued by the server:
    ASSERT_EQ(returned_responses.size(), 2);

    ASSERT_NE(returned_responses.at(0).path(),
              returned_responses.at(1).path());

    ASSERT_EQ(files_to_capture_set.count(returned_responses.at(0).path()), 1);
    ASSERT_EQ(files_to_capture_set.count(returned_responses.at(1).path()), 1);

    ASSERT_EQ(returned_responses.at(0).status().code(), grpc::StatusCode::OK);
    ASSERT_EQ(returned_responses.at(1).status().code(), grpc::StatusCode::OK);

    ASSERT_EQ(returned_responses.at(0).node_properties().mtime(), mtime);
    ASSERT_EQ(returned_responses.at(1).node_properties().mtime(), mtime);
}

TEST_F(ClientTestFixture, CaptureFilesErrorThrows)
{
    // The retry logic throws after running out of tries:
    EXPECT_CALL(*localCasClient.get(), CaptureFiles(_, _, _))
        .WillOnce(Return(
            grpc::Status(grpc::StatusCode::UNKNOWN, "Something went wrong.")));

    ASSERT_THROW(this->captureFiles({"/path/to/stage/file.txt"}, {}, false),
                 std::runtime_error);
}

TEST_F(ClientTestFixture, FetchTree)
{
    FetchTreeRequest request;
    EXPECT_CALL(*localCasClient.get(), FetchTree(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(grpc::Status::OK)));

    Digest digest;
    digest.set_hash("treeHash");
    digest.set_size_bytes(1234);

    ASSERT_NO_THROW(fetchTree(digest, false));

    EXPECT_EQ(request.instance_name(), this->getGrpcClient()->instanceName());
    EXPECT_EQ(request.root_digest(), digest);
    EXPECT_FALSE(request.fetch_file_blobs());
}

TEST_F(ClientTestFixture, FetchTreeWithFiles)
{
    FetchTreeRequest request;
    EXPECT_CALL(*localCasClient.get(), FetchTree(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(grpc::Status::OK)));

    Digest digest;
    digest.set_hash("treeHash");
    digest.set_size_bytes(1234);

    ASSERT_NO_THROW(fetchTree(digest, true));

    EXPECT_EQ(request.instance_name(), this->getGrpcClient()->instanceName());
    EXPECT_EQ(request.root_digest(), digest);
    EXPECT_TRUE(request.fetch_file_blobs());
}

TEST_F(ClientTestFixture, FetchTreeFails)
{
    FetchTreeRequest request;
    EXPECT_CALL(*localCasClient.get(), FetchTree(_, _, _))
        .Times(2)
        .WillRepeatedly(DoAll(SaveArg<1>(&request),
                              Return(grpc::Status(grpc::StatusCode::INTERNAL,
                                                  "Something went wrong."))));

    Digest digest;
    digest.set_hash("d");
    digest.set_size_bytes(1);

    ASSERT_THROW(fetchTree(digest, false), GrpcError);
    ASSERT_THROW(fetchTree(digest, true), GrpcError);
}

TEST_F(ClientTestFixture, FetchTreeFailsWithRetryableError)
{
    FetchTreeRequest request;
    EXPECT_CALL(*localCasClient.get(), FetchTree(_, _, _))
        .Times(4 * (getGrpcClient()->retryLimit() + 1))
        .WillRepeatedly(
            DoAll(SaveArg<1>(&request),
                  Return(grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                      "Something went wrong."))));

    Digest digest;
    digest.set_hash("d");
    digest.set_size_bytes(1);

    ASSERT_THROW(fetchTree(digest, false), GrpcError);
    ASSERT_THROW(fetchTree(digest, true), GrpcError);

    GrpcClient::RequestStats requestStats;
    ASSERT_THROW(fetchTree(digest, false, &requestStats), GrpcError);
    ASSERT_EQ(requestStats.d_grpcRetryCount, getGrpcClient()->retryLimit());

    ASSERT_THROW(fetchTree(digest, true, &requestStats), GrpcError);
    EXPECT_EQ(requestStats.d_grpcRetryCount,
              2 * getGrpcClient()->retryLimit());
}

class GetTreeFixture : public ClientTestFixture {
  protected:
    Digest d_digest;
    std::vector<buildboxcommon::Directory> d_directories;

    void prepareDigest()
    {
        /* Creates the following directory structure:
         *
         * ./
         *   src/
         *       build.sh*
         *       headers/
         *               file1.h
         *               file2.h
         *               file3.h
         *       cpp/
         *           file1.cpp
         *           file2.cpp
         *           file3.cpp
         *           symlink: file4.cpp --> file3.cpp
         */

        // ./src/headers
        Directory headers_directory;
        std::vector<std::string> headerFiles = {"file1.h", "file2.h",
                                                "file3.h"};
        for (const auto &file : headerFiles) {
            FileNode *fileNode = headers_directory.add_files();
            fileNode->set_name(file);
            fileNode->set_is_executable(false);
            fileNode->mutable_digest()->CopyFrom(
                DigestGenerator::hash(file + "_contents"));
        }
        const auto headers_directory_digest =
            DigestGenerator::hash(headers_directory);

        // ./src/cpp
        Directory cpp_directory;
        std::vector<std::string> cppFiles = {"file1.cpp", "file2.cpp",
                                             "file3.cpp"};
        for (const auto &file : cppFiles) {
            FileNode *fileNode = cpp_directory.add_files();
            fileNode->set_name(file);
            fileNode->set_is_executable(false);
            fileNode->mutable_digest()->CopyFrom(
                DigestGenerator::hash(file + "_contents"));
        }
        SymlinkNode *symNode = cpp_directory.add_symlinks();
        symNode->set_name("file4.cpp");
        symNode->set_target("file3.cpp");
        const auto cpp_directory_digest = DigestGenerator::hash(cpp_directory);

        // ./src
        Directory src_directory;
        DirectoryNode *headersNode = src_directory.add_directories();
        headersNode->set_name("headers");
        headersNode->mutable_digest()->CopyFrom(headers_directory_digest);
        DirectoryNode *cppNode = src_directory.add_directories();
        cppNode->set_name("cpp");
        cppNode->mutable_digest()->CopyFrom(cpp_directory_digest);
        FileNode *fileNode = src_directory.add_files();
        fileNode->set_name("build.sh");
        fileNode->set_is_executable(true);
        fileNode->mutable_digest()->CopyFrom(
            DigestGenerator::hash("build.sh_contents"));
        const auto src_directory_digest = DigestGenerator::hash(src_directory);

        // .
        Directory root_directory;
        DirectoryNode *srcNode = root_directory.add_directories();
        srcNode->set_name("src");
        srcNode->mutable_digest()->CopyFrom(src_directory_digest);

        const auto root_directory_serialized =
            root_directory.SerializeAsString();
        d_digest = DigestGenerator::hash(root_directory_serialized);

        d_directories.emplace_back(root_directory);
        d_directories.emplace_back(src_directory);
        d_directories.emplace_back(cpp_directory);
        d_directories.emplace_back(headers_directory);
    }

    GetTreeFixture() { prepareDigest(); }
};

TEST_F(GetTreeFixture, GetTreeSuccess)
{
    auto gettreereader = std::make_unique<MockGetTreeReader>();

    // prepare request
    GetTreeRequest request;
    request.set_instance_name("");
    request.mutable_root_digest()->CopyFrom(d_digest);

    EXPECT_CALL(*casClient, GetTreeRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(gettreereader.get())));

    // prepare expected response
    GetTreeResponse response;
    for (const auto &d : d_directories) {
        response.add_directories()->CopyFrom(d);
    }

    EXPECT_CALL(*gettreereader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*gettreereader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = gettreereader.release();
    (void)_;

    auto result = this->getTree(d_digest);
    ASSERT_EQ(result.map.size(), d_directories.size());
}

TEST_F(GetTreeFixture, GetTreeIncomplete)
{
    // Remove the last (headers) directory to test incomplete GetTree responses
    d_directories.pop_back();

    auto gettreereader = std::make_unique<MockGetTreeReader>();

    // prepare request
    GetTreeRequest request;
    request.set_instance_name("");
    request.mutable_root_digest()->CopyFrom(d_digest);

    EXPECT_CALL(*casClient, GetTreeRaw(_, _))
        .WillOnce(DoAll(SaveArg<1>(&request), Return(gettreereader.get())));

    // prepare expected response
    GetTreeResponse response;
    for (const auto &d : d_directories) {
        response.add_directories()->CopyFrom(d);
    }

    EXPECT_CALL(*gettreereader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*gettreereader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = gettreereader.release();
    (void)_;

    // Validation in `getTree()` should fail due to the incomplete response
    EXPECT_THROW(this->getTree(d_digest), std::runtime_error);
}

TEST_F(GetTreeFixture, GetTreeFail)
{
    auto gettreereader = std::make_unique<MockGetTreeReader>();

    EXPECT_CALL(*casClient, GetTreeRaw(_, _))
        .WillOnce(Return(gettreereader.get()));

    EXPECT_CALL(*gettreereader, Read(_))
        .WillOnce(Return(true))
        .WillOnce(Return(false));
    EXPECT_CALL(*gettreereader, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::StatusCode::NOT_FOUND,
                         "The root digest was not found in the local CAS.")));
    auto _ = gettreereader.release();
    (void)_;

    EXPECT_THROW(this->getTree(d_digest), std::runtime_error);
}

class UploadFileFixture : public ClientTestFixture {
    /**
     * Instantiates a tempfile with some data for use in upload tests.
     * Inherits from the pre-instantiated client fixture.
     */
  protected:
    UploadFileFixture()
    {
        if (write(tmpfile.fd(), content.c_str(), content.length()) !=
            static_cast<ssize_t>(content.length())) {
            throw "failed to upload file fixture";
        }
    }
};

TEST_F(UploadFileFixture, UploadFileTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    WriteRequest request;
    EXPECT_CALL(*writer, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(tmpfile.fd(), digest);
    EXPECT_EQ(request.resource_name().substr(0, client_instance_name.size()),
              client_instance_name);
}

TEST_F(UploadFileFixture, UploadFileCommittedSizeMismatch)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes() - 1);
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    WriteRequest request;
    EXPECT_CALL(*writer, Write(_, _))
        .WillOnce(DoAll(SaveArg<0>(&request), Return(true)));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    ASSERT_THROW(this->upload(tmpfile.fd(), digest), std::runtime_error);
}

TEST_F(UploadFileFixture, UploadLargeFileTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 3 * 1024 * 1024;
    std::string content = std::string(contentLength, 'f');

    // Overwrite the whole file
    lseek(tmpfile.fd(), 0, SEEK_SET);
    ASSERT_EQ(write(tmpfile.fd(), content.c_str(), contentLength),
              contentLength);

    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).Times(3).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(tmpfile.fd(), digest);
}

TEST_F(UploadFileFixture, UploadLargeFileRetryTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 3 * 1024 * 1024;
    std::string content = std::string(contentLength, 'f');

    // Overwrite the whole file
    lseek(tmpfile.fd(), 0, SEEK_SET);
    ASSERT_EQ(write(tmpfile.fd(), content.c_str(), contentLength),
              contentLength);

    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    auto failingWriter = std::make_unique<MockBytestreamWriter>();

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .Times(2)
        .WillOnce(Return(failingWriter.get()))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    // Fail first attempt with retryable error
    EXPECT_CALL(*failingWriter, Write(_, _)).WillOnce(Return(false));
    EXPECT_CALL(*failingWriter, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*failingWriter, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::UNAVAILABLE,
                                      "Something is wrong right now.")));
    auto _fw = failingWriter.release();
    (void)_fw;

    // Successful retry
    EXPECT_CALL(*writer, Write(_, _)).Times(3).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _w = writer.release();
    (void)_w;

    this->upload(tmpfile.fd(), digest);
}

TEST_F(UploadFileFixture, UploadExactFileTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 1024 * 1024;
    std::string content = std::string(contentLength, 'f');

    // Overwrite the whole file
    lseek(tmpfile.fd(), 0, SEEK_SET);
    ASSERT_EQ(write(tmpfile.fd(), content.c_str(), contentLength),
              contentLength);

    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(tmpfile.fd(), digest);
}

TEST_F(UploadFileFixture, UploadJustLargerThanExactFileTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 1024 * 1024 + 1;
    std::string content = std::string(contentLength, 'f');

    // Overwrite the whole file
    lseek(tmpfile.fd(), 0, SEEK_SET);
    ASSERT_EQ(write(tmpfile.fd(), content.c_str(), contentLength),
              contentLength);

    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).Times(2).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(tmpfile.fd(), digest);
}

TEST_F(UploadFileFixture, UploadJustSmallerThanExactFileTest)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    int contentLength = 1024 * 1024 - 1;
    std::string content = std::string(contentLength, 'f');

    // Overwrite the whole file
    lseek(tmpfile.fd(), 0, SEEK_SET);
    ASSERT_EQ(write(tmpfile.fd(), content.c_str(), contentLength),
              contentLength);

    digest.set_size_bytes(contentLength);
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    this->upload(tmpfile.fd(), digest);
}

TEST_F(UploadFileFixture, UploadFileReadFailure)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.release()));

    EXPECT_THROW(this->upload(-40, digest), std::system_error);
}

TEST_F(UploadFileFixture, UploadFileWriteFailure)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.get()));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(false));
    auto _ = writer.release();
    (void)_;

    EXPECT_THROW(this->upload(tmpfile.fd(), digest), std::runtime_error);
}

TEST_F(UploadFileFixture, UploadAlreadyExistingFile)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(false));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish())
        .WillOnce(Return(grpc::Status(grpc::Status::OK)));
    auto _ = writer.release();
    (void)_;

    ASSERT_NO_THROW(this->upload(tmpfile.fd(), digest));
}

TEST_F(UploadFileFixture, UploadFileDidntReturnOk)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    digest.set_size_bytes(static_cast<int64_t>(content.length()));
    digest.set_hash("fakehash");

    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(Return(writer.get()));

    EXPECT_CALL(*writer, Write(_, _)).WillOnce(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish())
        .WillOnce(Return(
            grpc::Status(grpc::FAILED_PRECONDITION, "failing for test")));
    auto _ = writer.release();
    (void)_;

    EXPECT_THROW(this->upload(tmpfile.fd(), digest), std::runtime_error);
}

class TransferDirectoryFixture : public ClientTestFixture {
    /**
     * Instantiates a tempfile with some data for use in upload tests.
     * Inherits from the pre-instantiated client fixture.
     */
  protected:
    /* directory/
     *   |-- file_a
     *   |-- subdirectory/
     *       |-- file_b
     */

    TransferDirectoryFixture()
        : ClientTestFixture(max_batch_size_bytes), directory(),
          subdirectory(directory.name(), "tmp-subdir"),
          file_a(directory.name(), "test-tmp-file"),
          file_b(subdirectory.name(), "test-tmp-file")
    {
        std::ofstream file_a_stream(file_a.name());
        file_a_stream << file_a_contents;
        file_a_stream.close();

        std::ofstream file_b_stream(file_b.name());
        file_b_stream << file_b_contents;
        file_b_stream.close();

        FileDescriptor dirfd(open(directory.name(), O_RDONLY | O_DIRECTORY));
        auto [nested_directory, fileMap] =
            Merklizer(false).makeNestedDirectory(dirfd.get());
        directory_file_map = // NOLINT
                             // (cppcoreguidelines-prefer-member-initializer)
            std::move(fileMap);

        directory_digest = nested_directory.to_digest(&directory_file_map);
        serialized_directory = directory_file_map.at(directory_digest);
    }

    // Setting a larger value so that we are allowed to upload complete
    // directories in a single batch
    static const int64_t max_batch_size_bytes = 1024UL * 1024UL;

    TemporaryDirectory directory, subdirectory;
    TemporaryFile file_a, file_b;

    const std::string file_a_contents = "Hello world!";
    const std::string file_b_contents = "This is some data...";

    const Digest file_a_digest = DigestGenerator::hash(file_a_contents);
    const Digest file_b_digest = DigestGenerator::hash(file_b_contents);

    digest_string_map directory_file_map;
    NestedDirectory nested_directory;
    std::string serialized_directory;
    Digest directory_digest;
};

TEST_F(TransferDirectoryFixture, UploadDirectory)
{
    // We expect the client to check if there are any blobs are missing to
    // avoid transferring those. For this test, we'll mock that all are missing
    // in the remote.

    // 1) FindMissingBlobs()
    FindMissingBlobsResponse missing_blobs_response;
    for (const auto &entry : directory_file_map) {
        auto missing_digest =
            missing_blobs_response.add_missing_blob_digests();
        missing_digest->CopyFrom(entry.first);
    }

    EXPECT_CALL(*casClient.get(), FindMissingBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(missing_blobs_response),
                        Return(grpc::Status::OK)));

    BatchUpdateBlobsRequest update_request;
    BatchUpdateBlobsResponse update_response;

    // 2) BatchUpdateBlobs()
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&update_request),
                        SetArgPointee<2>(update_response),
                        Return(grpc::Status::OK)));

    // We return success for all the updates:
    for (const auto &entry : directory_file_map) {
        const Digest file_digest = entry.first;
        auto response = update_response.add_responses();
        response->mutable_digest()->CopyFrom(file_digest);
        response->mutable_status()->set_code(grpc::StatusCode::OK);
    }

    Digest returned_directory_digest;
    this->uploadDirectory(std::string(directory.name()),
                          &returned_directory_digest);

    ASSERT_EQ(returned_directory_digest, directory_digest);

    // All the data was written to the remote:
    ASSERT_EQ(update_request.requests_size(), directory_file_map.size());
    for (const auto &entry : update_request.requests()) {
        ASSERT_EQ(directory_file_map.count(entry.digest()), 1);
    }
}

TEST_F(TransferDirectoryFixture, UploadDirectoryLargeFile)
{
    auto writer = std::make_unique<MockBytestreamWriter>();

    const int contentLength = 3 * 1024 * 1024;
    const std::string content = std::string(contentLength, 'f');
    const auto digest = DigestGenerator::hash(content);
    std::ofstream largeFile(std::string(directory.name()) + "/large-file");
    largeFile << content;
    largeFile.close();

    // Update directory file map and digest with the large file
    FileDescriptor dirfd(open(directory.name(), O_RDONLY | O_DIRECTORY));
    auto [nested_large_directory, large_directory_file_map] =
        Merklizer(false).makeNestedDirectory(dirfd.get());
    const auto large_directory_digest =
        nested_large_directory.to_digest(&large_directory_file_map);

    // We expect the client to check if there are any blobs are missing to
    // avoid transferring those. For this test, we'll mock that all are missing
    // in the remote.

    // 1) FindMissingBlobs()
    FindMissingBlobsResponse missing_blobs_response;
    for (const auto &entry : large_directory_file_map) {
        auto missing_digest =
            missing_blobs_response.add_missing_blob_digests();
        missing_digest->CopyFrom(entry.first);
    }

    EXPECT_CALL(*casClient.get(), FindMissingBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(missing_blobs_response),
                        Return(grpc::Status::OK)));

    BatchUpdateBlobsRequest update_request;
    BatchUpdateBlobsResponse update_response;

    // 2) BatchUpdateBlobs() for the small files
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&update_request),
                        SetArgPointee<2>(update_response),
                        Return(grpc::Status::OK)));

    // We return success for all the updates:
    for (const auto &entry : large_directory_file_map) {
        const Digest file_digest = entry.first;
        if (file_digest.size_bytes() <= MAX_BATCH_SIZE_BYTES) {
            auto response = update_response.add_responses();
            response->mutable_digest()->CopyFrom(file_digest);
            response->mutable_status()->set_code(grpc::StatusCode::OK);
        }
    }

    // 3) Write() for the large file
    writeResponse.set_committed_size(digest.size_bytes());
    EXPECT_CALL(*bytestreamClient, WriteRaw(_, _))
        .WillOnce(
            DoAll(SetArgPointee<1>(writeResponse), Return(writer.get())));

    EXPECT_CALL(*writer, Write(_, _)).Times(3).WillRepeatedly(Return(true));
    EXPECT_CALL(*writer, WritesDone()).WillOnce(Return(true));
    EXPECT_CALL(*writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = writer.release();
    (void)_;

    Digest returned_directory_digest;
    this->uploadDirectory(std::string(directory.name()),
                          &returned_directory_digest);
    ASSERT_EQ(returned_directory_digest, large_directory_digest);
}

TEST_F(TransferDirectoryFixture, UploadDirectoryNoMissingBlobs)
{
    // In this test the remote reports that no blobs are missing, so no upload
    // needs to take place.

    FindMissingBlobsResponse missing_blobs_response;

    EXPECT_CALL(*casClient.get(), FindMissingBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(missing_blobs_response),
                        Return(grpc::Status::OK)));

    Digest returned_directory_digest;
    this->uploadDirectory(std::string(directory.name()),
                          &returned_directory_digest);

    ASSERT_EQ(returned_directory_digest, directory_digest);
}

TEST_F(TransferDirectoryFixture, UploadDirectoryWritesTree)
{
    /* directory/
     *   |-- file_a
     *   |-- subdirectory/
     *       |-- file_b
     */

    FindMissingBlobsResponse missing_blobs_response;
    // The remote reports that no blobs are missing, so no upload
    // needs to take place.
    EXPECT_CALL(*casClient.get(), FindMissingBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(missing_blobs_response),
                        Return(grpc::Status::OK)));

    Digest returned_directory_digest;
    Tree tree;
    this->uploadDirectory(std::string(directory.name()),
                          &returned_directory_digest, &tree);

    // `uploadDirectory()` wrote the `Tree` object that it generated before the
    // transfer, it should match the data:

    // directory/:
    ASSERT_EQ(tree.root().directories_size(), 1); // subdirectory/
    ASSERT_EQ(tree.children_size(), 1);

    ASSERT_EQ(tree.root().files_size(), 1); // fila_a
    ASSERT_EQ(tree.root().symlinks_size(), 0);

    // subdirectory/
    ASSERT_EQ(tree.children(0).files_size(), 1); // file_b
    ASSERT_EQ(tree.children(0).directories_size(), 0);
    ASSERT_EQ(tree.children(0).symlinks_size(), 0);

    ASSERT_EQ(returned_directory_digest, directory_digest);
}

TEST_F(ClientTestFixture, DownloadDirectoryTestActualDownload)
{
    TemporaryDirectory capture_dir;
    TemporaryDirectory write_dir;

    const mode_t mode = PERMISSION_RWXUSR;
    const std::chrono::system_clock::time_point mtime{};

    auto file_in_capture_dir = std::string(capture_dir.name()) + "/file1";
    auto symlink_in_capture_dir =
        std::string(capture_dir.name()) + "/" + "symlink1";
    auto symlink_in_write_dir =
        std::string(write_dir.name()) + "/" + "symlink1";
    auto file_in_write_dir = std::string(write_dir.name()) + "/file1";

    // create the file in the tempdir: capture_dir
    buildboxcommontest::TestUtils::touchFile(file_in_capture_dir.c_str());
    // make a symlink to the file
    EXPECT_TRUE(symlink(file_in_capture_dir.c_str(),
                        symlink_in_capture_dir.c_str()) == 0);
    // Set properties of capture_dir
    ASSERT_EQ(0, chmod(capture_dir.name(), mode));
    FileUtils::setFileMtime(capture_dir.name(), mtime);

    FileDescriptor dirfd(open(capture_dir.name(), O_RDONLY | O_DIRECTORY));
    auto result = Merklizer(false, {"mtime", "unix_mode"})
                      .makeNestedDirectory(dirfd.get());
    auto &nested_directory = result.first;
    digest_string_map &directory_file_map = result.second;
    Digest directory_digest = nested_directory.to_digest(&directory_file_map);

    // this callback, if called, will create the file in the write directory.
    download_callback_t download_blobs =
        [&](const std::vector<Digest> &file_digests,
            const OutputMap &outputs) {
            buildboxcommontest::TestUtils::touchFile(
                file_in_write_dir.c_str());
        };

    // return the Directory Node from the map created in make_nesteddirectory
    return_directory_callback_t download_directory =
        [&](const Digest &digest) {
            Directory dir;
            const auto dir_string = directory_file_map.at(directory_digest);
            dir.ParseFromString(dir_string);
            return dir;
        };

    GrpcClient::RequestStats requestStats;
    this->downloadDirectory(directory_digest, write_dir.name(), download_blobs,
                            download_directory, &requestStats);

    // verify that write_dir has the same contents as capture_dir
    ASSERT_FALSE(FileUtils::directoryIsEmpty(write_dir.name()));
    ASSERT_TRUE(FileUtils::isDirectory(write_dir.name()));
    ASSERT_TRUE(FileUtils::isRegularFile(file_in_write_dir.c_str()));
    ASSERT_TRUE(FileUtils::isSymlink(symlink_in_write_dir.c_str()));
    EXPECT_EQ(mode, FileUtils::getUnixMode(write_dir.name()) & 0777);
    EXPECT_EQ(
        mtime.time_since_epoch().count(),
        FileUtils::getFileMtime(write_dir.name()).time_since_epoch().count());

    ASSERT_EQ(requestStats.d_grpcRetryCount, 0);
}

TEST_F(TransferDirectoryFixture, DownloadDirectoryMissingDigestThrows)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    Digest digest;
    digest.set_hash("ThisDoesNotExist");
    digest.set_size_bytes(1234);

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));

    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    EXPECT_CALL(*reader, Finish())
        .WillOnce(Return(grpc::Status(grpc::StatusCode::NOT_FOUND,
                                      "Blob not found in CAS")));
    auto _ = reader.release();
    (void)_;

    TemporaryDirectory output_dir;
    ASSERT_THROW(this->downloadDirectory(digest, output_dir.name()),
                 std::runtime_error);
}

TEST_F(TransferDirectoryFixture, DownloadDirectoryRetryableErrorThrows)
{
    Digest digest;
    digest.set_hash("ThisDoesNotExist");
    digest.set_size_bytes(1234);

    auto reader1 = std::make_unique<MockBytestreamReader>();
    auto reader2 = std::make_unique<MockBytestreamReader>();

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader1.get()))
        .WillOnce(Return(reader2.get()));

    const auto errorStatus =
        grpc::Status(grpc::StatusCode::UNAVAILABLE, "Temporary issue");

    EXPECT_CALL(*reader1, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader1, Finish()).WillOnce(Return(errorStatus));
    auto _1 = reader1.release();
    (void)_1;

    EXPECT_CALL(*reader2, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader2, Finish()).WillOnce(Return(errorStatus));
    auto _2 = reader2.release();
    (void)_2;

    TemporaryDirectory output_dir;
    GrpcClient::RequestStats requestStats;
    ASSERT_THROW(
        this->downloadDirectory(digest, output_dir.name(), &requestStats),
        std::runtime_error);

    ASSERT_EQ(requestStats.d_grpcRetryCount, getGrpcClient()->retryLimit());
}

TEST_F(TransferDirectoryFixture, StageDirectory)
{
    auto reader_writer = std::make_unique<MockGetTreeReaderWriter>();

    EXPECT_CALL(*localCasClient.get(), StageTreeRaw(_))
        .WillOnce(Return(reader_writer.get()));

    // The client will issue 2 requests: the actual `StageTreeRequest` and an
    // empty message to indicate to the server that it can clean up.
    EXPECT_CALL(*reader_writer, Write(_, _))
        .Times(2)
        .WillRepeatedly(Return(true));

    StageTreeResponse response;
    response.set_path("/tmp/stage/");
    EXPECT_CALL(*reader_writer, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(response), Return(true)));

    {
        std::unique_ptr<StagedDirectory> staged_dir =
            this->stage(directory_digest, StageTreeRequest_StagingMode_DEFAULT,
                        "/tmp/stage");
        ASSERT_EQ(staged_dir->path(), "/tmp/stage/");

        // The StagedDirectory destructor automatically sends the second empty
        // message to the server.
        EXPECT_CALL(*reader_writer, WritesDone()).WillOnce(Return(true));
    }
    auto _ = reader_writer.release();
    (void)_;
}

TEST_F(TransferDirectoryFixture, StageDirectoryThrowsOnError)
{
    auto reader_writer = std::make_unique<MockGetTreeReaderWriter>();

    EXPECT_CALL(*localCasClient.get(), StageTreeRaw(_))
        .WillOnce(Return(reader_writer.get()));

    EXPECT_CALL(*reader_writer, Write(_, _)).WillOnce(Return(false));

    EXPECT_CALL(*reader_writer, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader_writer.release();
    (void)_;

    EXPECT_THROW(this->stage(directory_digest,
                             StageTreeRequest_StagingMode_DEFAULT,
                             "/tmp/stage"),
                 std::runtime_error);
}

TEST_F(TransferDirectoryFixture, UploadBlobsWithFdAndPath)
{
    FileDescriptor fd(open(file_a.name(), O_RDONLY));
    ASSERT_GE(fd.get(), 0);

    auto request_fd =
        CASClient::UploadRequest::from_fd(file_a_digest, fd.get());
    auto request_path =
        CASClient::UploadRequest::from_path(file_b_digest, file_b.name());
    std::vector<CASClient::UploadRequest> requests = {request_fd,
                                                      request_path};

    BatchUpdateBlobsResponse response;
    response.add_responses()->mutable_digest()->CopyFrom(file_a_digest);

    BatchUpdateBlobsRequest batchRequest;
    EXPECT_CALL(*casClient.get(), BatchUpdateBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&batchRequest), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));

    const auto failed_uploads = this->uploadBlobs(requests);
    EXPECT_TRUE(failed_uploads.empty());

    EXPECT_EQ(batchRequest.requests_size(), 2);
    EXPECT_EQ(batchRequest.requests(0).digest(), file_a_digest);
    EXPECT_EQ(batchRequest.requests(0).data(), file_a_contents);
    EXPECT_EQ(batchRequest.requests(1).digest(), file_b_digest);
    EXPECT_EQ(batchRequest.requests(1).data(), file_b_contents);
}

class DownloadBlobsFixture : public ClientTestFixture {
  protected:
    DownloadBlobsFixture() {}
};

TEST_F(DownloadBlobsFixture, FileTooLargeToBatchDownload)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    // Expecting it to fall back to a bytestream Read():
    const auto data = std::string(2 * MAX_BATCH_SIZE_BYTES, '-');
    const Digest digest = DigestGenerator::hash(data);

    const std::vector<Digest> requests = {digest};

    readResponse.set_data(data);

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));

    auto write_blob = [&](const std::string &downloaded_hash,
                          const std::string &downloaded_data) {
        ASSERT_EQ(downloaded_hash, digest.hash());
        ASSERT_EQ(downloaded_data, data);
    };

    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    ASSERT_NO_THROW(this->downloadBlobs(requests, write_blob, -1, nullptr));
}

TEST_F(DownloadBlobsFixture, DownloadBlobs)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    const std::vector<std::string> payload = {
        "a", "b", std::string(3 * MAX_BATCH_SIZE_BYTES, 'x'), "c"};
    std::vector<std::string> hashes;

    // Creating list of requests...
    std::vector<Digest> requests;
    for (unsigned i = 0; i < payload.size(); i++) {
        const Digest digest = DigestGenerator::hash(payload[i]);
        hashes.push_back(digest.hash());
        requests.push_back(digest);
    }
    ASSERT_EQ(requests.size(), payload.size());

    // We expect digests {0, 1, 3} to be requested with BatchReadBlobs().
    BatchReadBlobsResponse response;
    for (unsigned i = 0; i < payload.size(); i++) {
        if (i == 2)
            continue;
        auto entry = response.add_responses();
        entry->set_data(payload[i]);
        entry->mutable_digest()->CopyFrom(requests[i]);
    }

    BatchReadBlobsRequest request;
    EXPECT_CALL(*casClient.get(), BatchReadBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));

    // And digest in index 2 with the Bytestream API:
    readResponse.set_data(payload[2]);
    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    // We will write the output results into a map indexed by hash:
    std::unordered_map<std::string, std::string> downloads;
    auto write_blob = [&](const std::string &hash, const std::string &data) {
        downloads[hash] = data;
    };

    CASClient::DownloadResults download_results;
    ASSERT_NO_THROW(download_results = this->downloadBlobs(
                        requests, write_blob, -1, nullptr));

    // Client sends the correct instance name:
    EXPECT_EQ(request.instance_name(), client_instance_name);

    // We get all the data, and it's correct for each requested digest:
    ASSERT_EQ(downloads.size(), requests.size());
    ASSERT_EQ(downloads.at(hashes[0]), payload[0]);
    ASSERT_EQ(downloads.at(hashes[1]), payload[1]);
    ASSERT_EQ(downloads.at(hashes[2]), payload[2]);
    ASSERT_EQ(downloads.at(hashes[3]), payload[3]);

    // The returned dictionary has whether or not each digest was successfully
    // fetched:
    ASSERT_EQ(download_results.size(), 4);
    for (const auto &digest : requests) {
        const auto result = std::find_if(
            download_results.cbegin(), download_results.cend(),
            [&digest](const DownloadResult &r) { return r.first == digest; });

        ASSERT_NE(result, download_results.cend());
        ASSERT_EQ(result->second.code(), grpc::StatusCode::OK);
    }
}

TEST_F(DownloadBlobsFixture, DownloadBlobsThreaded)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    ThreadPool pool(2, "");
    const std::vector<std::string> payload = {
        "a", "b", std::string(3 * MAX_BATCH_SIZE_BYTES, 'x'), "c"};
    std::vector<std::string> hashes;

    // Creating list of requests...
    std::vector<Digest> requests;
    for (unsigned i = 0; i < payload.size(); i++) {
        const Digest digest = DigestGenerator::hash(payload[i]);
        hashes.push_back(digest.hash());
        requests.push_back(digest);
    }
    ASSERT_EQ(requests.size(), payload.size());

    // We expect digests {0, 1, 3} to be requested with BatchReadBlobs().
    BatchReadBlobsResponse response;
    for (unsigned i = 0; i < payload.size(); i++) {
        if (i == 2)
            continue;
        auto entry = response.add_responses();
        entry->set_data(payload[i]);
        entry->mutable_digest()->CopyFrom(requests[i]);
    }

    BatchReadBlobsRequest request;
    EXPECT_CALL(*casClient.get(), BatchReadBlobs(_, _, _))
        .WillOnce(DoAll(SaveArg<1>(&request), SetArgPointee<2>(response),
                        Return(grpc::Status::OK)));

    // And digest in index 2 with the Bytestream API:
    readResponse.set_data(payload[2]);
    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    // We will write the output results into a map indexed by hash:
    std::unordered_map<std::string, std::string> downloads;
    auto write_blob = [&](const std::string &hash, const std::string &data) {
        downloads[hash] = data;
    };

    CASClient::DownloadResults download_results;
    ASSERT_NO_THROW(download_results =
                        this->downloadBlobs(requests, write_blob, -1, &pool));

    // Client sends the correct instance name:
    EXPECT_EQ(request.instance_name(), client_instance_name);

    // We get all the data, and it's correct for each requested digest:
    ASSERT_EQ(downloads.size(), requests.size());
    ASSERT_EQ(downloads.at(hashes[0]), payload[0]);
    ASSERT_EQ(downloads.at(hashes[1]), payload[1]);
    ASSERT_EQ(downloads.at(hashes[2]), payload[2]);
    ASSERT_EQ(downloads.at(hashes[3]), payload[3]);

    // The returned dictionary has whether or not each digest was successfully
    // fetched:
    ASSERT_EQ(download_results.size(), 4);
    for (const auto &digest : requests) {
        const auto result = std::find_if(
            download_results.cbegin(), download_results.cend(),
            [&digest](const DownloadResult &r) { return r.first == digest; });

        ASSERT_NE(result, download_results.cend());
        ASSERT_EQ(result->second.code(), grpc::StatusCode::OK);
    }
}

TEST_F(DownloadBlobsFixture, DownloadBlobsBatchWithMissingBlob)
{
    const Digest existing_digest = DigestGenerator::hash("existing-blob");
    const Digest non_existing_digest = DigestGenerator::hash("missing-blob");

    BatchReadBlobsResponse response;
    auto entry1 = response.add_responses();
    entry1->mutable_digest()->CopyFrom(non_existing_digest);
    entry1->mutable_status()->set_code(grpc::StatusCode::NOT_FOUND);
    entry1->mutable_status()->set_message("Digest not found in CAS.");

    auto entry2 = response.add_responses();
    entry2->mutable_digest()->CopyFrom(existing_digest);
    entry2->mutable_status()->set_code(grpc::StatusCode::OK);
    entry2->set_data("existing-blob");

    EXPECT_CALL(*casClient.get(), BatchReadBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(response), Return(grpc::Status::OK)));

    unsigned int written_blobs = 0;
    const auto write_blob = [&](const std::string &hash, const std::string &) {
        ASSERT_EQ(hash, existing_digest.hash());
        written_blobs++;
    };

    std::vector<Digest> requests = {existing_digest, non_existing_digest};
    CASClient::DownloadResults download_results;
    ASSERT_NO_THROW(download_results = this->downloadBlobs(
                        requests, write_blob, -1, nullptr));
    ASSERT_EQ(written_blobs, 1);

    ASSERT_EQ(download_results.size(), 2);

    for (const auto &entry : download_results) {
        const Digest digest = entry.first;
        const google::rpc::Status status = entry.second;

        if (digest == existing_digest) {
            ASSERT_EQ(status.code(), grpc::StatusCode::OK);
        }
        else if (digest == non_existing_digest) {
            ASSERT_EQ(status.code(), grpc::StatusCode::NOT_FOUND);
        }
        else {
            FAIL() << "Unexpected digest in response: [" << digest
                   << "] was not requested.";
        }
    }
}

TEST_F(DownloadBlobsFixture, DownloadBlobsBatchReadBlobsFails)
{
    Digest digest;
    digest.set_hash("hash0");
    digest.set_size_bytes(1);

    unsigned int written_blobs = 0;
    const auto write_blob = [&](const std::string &, const std::string &) {
        written_blobs++;
    };

    auto result = BatchReadBlobsResponse();
    auto *response = result.add_responses();
    *response->mutable_digest() = digest;
    auto *status = response->mutable_status();
    status->set_code(grpc::StatusCode::NOT_FOUND);
    status->set_message("Digest not found");

    EXPECT_CALL(*casClient.get(), BatchReadBlobs(_, _, _))
        .WillOnce(DoAll(SetArgPointee<2>(result), Return(grpc::Status::OK)));

    CASClient::DownloadResults download_results;
    ASSERT_NO_THROW(download_results = this->downloadBlobs(
                        {digest}, write_blob, -1, nullptr));

    ASSERT_EQ(written_blobs, 0);

    // The returned vector contains that the requested digest failed
    // to be downloaded:
    ASSERT_EQ(download_results.size(), 1);
    ASSERT_EQ(download_results.front().first, digest);
    ASSERT_EQ(download_results.front().second.code(), status->code());
}

TEST_F(DownloadBlobsFixture, DownloadBlobsHelperFails)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    Digest digest;
    digest.set_hash("hash0");
    digest.set_size_bytes(3 * MAX_BATCH_SIZE_BYTES);

    unsigned int written_blobs = 0;
    const auto write_blob = [&](const std::string &, const std::string &) {
        written_blobs++;
    };

    const auto errorStatus =
        grpc::Status(grpc::StatusCode::NOT_FOUND, "Digest not found in CAS.");
    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_)).WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(errorStatus));
    auto _ = reader.release();
    (void)_;

    CASClient::DownloadResults download_results;
    ASSERT_NO_THROW(download_results = this->downloadBlobs(
                        {digest}, write_blob, -1, nullptr));

    ASSERT_EQ(written_blobs, 0);

    // The returned vector contains that the requested digest failed
    // to be downloaded:
    ASSERT_EQ(download_results.size(), 1);
    ASSERT_EQ(download_results.front().first, digest);
    ASSERT_EQ(download_results.front().second.code(),
              errorStatus.error_code());
}

TEST_F(DownloadBlobsFixture, DownloadBlobsResultSuccessfulStatusAndData)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    // Test the public `downloadBlobs()` method to check that the returned map
    // is correct.
    const auto data = std::string(MAX_BATCH_SIZE_BYTES + 1, 'A');
    readResponse.set_data(data);

    const Digest digest = DigestGenerator::hash(data);

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    const CASClient::DownloadBlobsResult download_results =
        this->downloadBlobs({digest}, nullptr);
    ASSERT_EQ(download_results.count(digest.hash()), 1);
    ASSERT_EQ(download_results.size(), 1);

    const auto &result = download_results.at(digest.hash());
    const auto &result_status = result.first;
    const auto &result_data = result.second;

    ASSERT_EQ(result_status.code(), grpc::StatusCode::OK);
    ASSERT_EQ(result_data, data);
}

TEST_F(DownloadBlobsFixture, DownloadBlobsResultErrorCode)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    // Test the public `downloadBlobs()` method to check that the returned map
    // is correct.
    Digest digest;
    digest.set_hash("hash0");
    digest.set_size_bytes(3 * MAX_BATCH_SIZE_BYTES);

    const auto errorStatus =
        grpc::Status(grpc::StatusCode::NOT_FOUND, "Digest not found in CAS.");
    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_)).WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(errorStatus));
    auto _ = reader.release();
    (void)_;

    CASClient::DownloadBlobsResult download_results;
    ASSERT_NO_THROW(download_results = this->downloadBlobs({digest}, nullptr));

    ASSERT_EQ(download_results.size(), 1);
    ASSERT_EQ(download_results.count(digest.hash()), 1);

    const auto result = download_results.at(digest.hash());
    ASSERT_EQ(result.first.code(), errorStatus.error_code());
    ASSERT_EQ(result.second, "");
}

TEST_F(DownloadBlobsFixture,
       DownloadBlobsToDirectoryResultSuccessfulStatusAndPath)
{
    auto reader = std::make_unique<MockBytestreamReader>();

    // Test the public `downloadBlobsToDirectory()` method to check that the
    // returned map is correct.
    const auto data = std::string(MAX_BATCH_SIZE_BYTES + 1, 'A');
    readResponse.set_data(data);

    const Digest digest = DigestGenerator::hash(data);

    EXPECT_CALL(*bytestreamClient, ReadRaw(_, _))
        .WillOnce(Return(reader.get()));
    EXPECT_CALL(*reader, Read(_))
        .WillOnce(DoAll(SetArgPointee<0>(readResponse), Return(true)))
        .WillOnce(Return(false));
    EXPECT_CALL(*reader, Finish()).WillOnce(Return(grpc::Status::OK));
    auto _ = reader.release();
    (void)_;

    TemporaryDirectory directory;

    const CASClient::DownloadBlobsResult download_results =
        this->downloadBlobsToDirectory({digest}, directory.fd());
    ASSERT_EQ(download_results.count(digest.hash()), 1);
    ASSERT_EQ(download_results.size(), 1);

    const auto &result = download_results.at(digest.hash());
    const auto &result_status = result.first;
    const auto result_path = directory.strname() + "/" + result.second;

    const auto result_data = FileUtils::getFileContents(result_path.c_str());

    ASSERT_EQ(result_status.code(), grpc::StatusCode::OK);
    ASSERT_EQ(result_data, data);
}
