/*
 * 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.
 */

#ifndef INCLUDED_BUILDBOXCASD_CASPROXY_H
#define INCLUDED_BUILDBOXCASD_CASPROXY_H

#include <buildboxcasd_metricnames.h>
#include <buildboxcasd_server.h>
#include <buildboxcommon_digestgenerator.h>
#include <buildboxcommonmetrics_countingmetricutil.h>
#include <buildboxcommonmetrics_distributionmetric.h>
#include <buildboxcommonmetrics_metriccollectorfactory.h>
#include <buildboxcommonmetrics_testingutils.h>

#include <thread>

#include <fcntl.h>
#include <gmock/gmock.h>

#include <grpcpp/client_context.h>
#include <grpcpp/create_channel.h>
#include <grpcpp/impl/codegen/call.h>
#include <grpcpp/server_context.h>
#include <grpcpp/support/sync_stream.h>

#include <gtest/gtest.h>

#include <buildboxcommon_temporarydirectory.h>
#include <buildboxcommonmetrics_countingmetricvalue.h>

#include <buildboxcasd_casinstance.h>
#include <buildboxcasd_proxymodefixture.h>

using namespace buildboxcasd;
using buildboxcasd::MetricNames;
using buildboxcommon::DigestGenerator;

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

namespace {
template <typename ValueType>
std::unordered_map<std::string, ValueType> getMetricSnapshot()
{
    buildboxcommonmetrics::MetricCollector<ValueType> *collector =
        buildboxcommonmetrics::MetricCollectorFactory::getCollector<
            ValueType>();
    return collector->getSnapshot();
}
template <typename ValueType>
using MetricVector = std::vector<std::pair<std::string, ValueType>>;

template <typename ValueType>
void validateMetrics(
    const std::vector<std::pair<std::string, ValueType>> expected_metrics)
{
    const auto actual_metrics = getMetricSnapshot<ValueType>();
    for (const auto &expected : expected_metrics) {
        const auto actual = actual_metrics.find(expected.first);
        if (actual == actual_metrics.end()) {
            FAIL() << "Expected metric " << expected.first
                   << " was not collected!";
        }
        else if (actual->second != expected.second) {
            FAIL() << "Metric " << expected.first
                   << " expected=" << expected.second.value()
                   << " actual=" << actual->second.value();
        }
    }
}
} // namespace

class CasProxyFixture : public CasProxyModeFixture {
  protected:
    CasProxyFixture()
    {
        // Pre-loading some content in the local and remote storages:
        local_storage->writeBlob(DigestGenerator::hash("data1"), "data1");
        local_storage->writeBlob(DigestGenerator::hash("data2"), "data2");

        remote_storage->writeBlob(DigestGenerator::hash("remotedata3"),
                                  "remotedata3");
        remote_storage->writeBlob(DigestGenerator::hash("remotedata4"),
                                  "remotedata4");

        // Clear out any metrics collected from the priming writeBlob calls
        buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    }

    void prepareTree(Digest *root_digest_out,
                     Digest *subdirectory_digest_out = nullptr)
    {
        /* Creates the following directory structure:
         *
         * root/                <-- `root_directory_digest_out`
         *  |-- file1.sh*
         *  |-- subdir1/        <-- `subdirectory_digest_out` (optional)
         *           |-- file2.c
         *
         * And returns pointers to the root directory and to the `subdir1/`
         * digests.
         *
         * If `subdirectory_digest_out` is a null pointer, it does not store
         * that subdirectory in the CAS.
         */

        Directory root_directory;

        // Adding a file to it:
        FileNode *f1 = root_directory.add_files();
        f1->set_name("file1.sh");
        f1->set_is_executable(true);
        f1->mutable_digest()->CopyFrom(DigestGenerator::hash("file1Contents"));

        // Creating `subdir1/`:
        Directory subdirectory;

        // Adding `subdir1/file2.c`:
        FileNode *f2 = subdirectory.add_files();
        f2->set_name("file2.c");
        f2->mutable_digest()->CopyFrom(DigestGenerator::hash("file2Contents"));

        // Adding `subdir1/` under `dirA/`:
        DirectoryNode *d1 = root_directory.add_directories();
        d1->set_name("subdir1");

        const auto serialized_subdirectory = subdirectory.SerializeAsString();

        const auto subdirectory_digest =
            DigestGenerator::hash(serialized_subdirectory);
        d1->mutable_digest()->CopyFrom(subdirectory_digest);

        const auto serialized_root_directory =
            root_directory.SerializeAsString();
        const auto root_directory_digest =
            DigestGenerator::hash(serialized_root_directory);

        // Storing the two Directory protos in the remote CAS:
        remote_storage->writeBlob(root_directory_digest,
                                  serialized_root_directory);

        if (subdirectory_digest_out != nullptr) {
            remote_storage->writeBlob(subdirectory_digest,
                                      serialized_subdirectory);
            *subdirectory_digest_out = subdirectory_digest;
        }

        *root_digest_out = root_directory_digest;
    }
};

TEST_P(CasProxyFixture, FindMissingBlobs)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    FindMissingBlobsRequest request;
    request.set_instance_name(instance_name);

    const auto local_digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(local_digest));

    const auto remote_digest = DigestGenerator::hash("remotedata4");
    ASSERT_TRUE(remote_storage->hasBlob(remote_digest));

    const auto missing_digest = DigestGenerator::hash("non-existent");
    ASSERT_FALSE(local_storage->hasBlob(missing_digest));
    ASSERT_FALSE(remote_storage->hasBlob(missing_digest));

    const auto r1 = request.add_blob_digests();
    r1->CopyFrom(local_digest);
    const auto r2 = request.add_blob_digests();
    r2->CopyFrom(remote_digest);
    const auto r3 = request.add_blob_digests();
    r3->CopyFrom(missing_digest);

    FindMissingBlobsResponse response;
    cas_servicer->FindMissingBlobs(&server_context, &request, &response);

    ASSERT_EQ(response.missing_blob_digests().size(), 1);
    ASSERT_EQ(response.missing_blob_digests(0).size_bytes(),
              missing_digest.size_bytes());
    ASSERT_EQ(response.missing_blob_digests(0).hash(), missing_digest.hash());

    // The blob missing in the remote but available locally was implicitly
    // uploaded:
    ASSERT_TRUE(remote_storage->hasBlob(local_digest));

    // and thus there was a local read and write to the remote from when we
    // pushed the blob up.
    const MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_FIND_MISSING,
         buildboxcommonmetrics::CountingMetricValue(3)},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(
             local_digest.size_bytes())},
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
         buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL,
         buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)},
    };

    validateMetrics<buildboxcommonmetrics::CountingMetricValue>(expected);

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::allCollectedByNameWithValues<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>({
            {MetricNames::DISTRIBUTION_NAME_CAS_FIND_MISSING_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 local_digest.size_bytes())},
            {MetricNames::DISTRIBUTION_NAME_CAS_FIND_MISSING_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 remote_digest.size_bytes())},
            {MetricNames::DISTRIBUTION_NAME_CAS_FIND_MISSING_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 missing_digest.size_bytes())},
        }));
}

TEST_P(CasProxyFixture, BatchUpdateBlobs)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();

    BatchUpdateBlobsRequest request;
    request.set_instance_name(instance_name);

    const auto data4 = "data4";
    const auto data5 = "data5";

    const Digest d4 = DigestGenerator::hash(data4);
    const Digest d5 = DigestGenerator::hash(data5);

    auto entry1 = request.add_requests();
    entry1->mutable_digest()->CopyFrom(d4);
    entry1->set_data(data4);

    auto entry2 = request.add_requests();
    entry2->mutable_digest()->CopyFrom(d5);
    entry2->set_data(data5);

    BatchUpdateBlobsResponse response;
    cas_servicer->BatchUpdateBlobs(&server_context, &request, &response);

    // Response is valid:
    ASSERT_EQ(response.responses().size(), 2);

    std::set<std::string> digests;
    for (const auto &r : response.responses()) {
        ASSERT_EQ(r.status().code(), grpc::StatusCode::OK);
        digests.insert(r.digest().hash());
    }

    ASSERT_EQ(digests.count(d4.hash()), 1);
    ASSERT_EQ(digests.count(d5.hash()), 1);

    // And data was stored:
    ASSERT_TRUE(local_storage->hasBlob(d4));
    ASSERT_TRUE(local_storage->hasBlob(d5));

    ASSERT_TRUE(remote_storage->hasBlob(d4));
    ASSERT_TRUE(remote_storage->hasBlob(d5));

    // Verify

    // each blob written to the remote cache, and both local caches are updated
    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_BATCH_UPDATE,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(d4.size_bytes() +
                                                    d5.size_bytes())},
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(4)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(2 * d4.size_bytes() +
                                                    2 * d5.size_bytes())},
    };
    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BatchUpdateBlobsWithUnreachableRemote)
{
    stopRemoteServer();

    BatchUpdateBlobsRequest request;
    request.set_instance_name(instance_name);

    auto entry = request.add_requests();
    const auto data = "data1";
    const auto digest = DigestGenerator::hash(data);
    entry->mutable_digest()->CopyFrom(digest);
    entry->set_data(data);

    BatchUpdateBlobsResponse response;
    grpc::Status status;
    ASSERT_NO_THROW(status = cas_servicer->BatchUpdateBlobs(
                        &server_context, &request, &response));

    ASSERT_EQ(status.error_code(), grpc::StatusCode::OK);

    ASSERT_EQ(response.responses_size(), 1);
    ASSERT_EQ(response.responses(0).digest(), digest);
    ASSERT_NE(response.responses(0).status().code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyFixture, BatchReadBlobs)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    // We'll request two blobs that we know are in the CAS:
    const auto digest1 = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest1));

    const auto digest2 = DigestGenerator::hash("remotedata3");
    ASSERT_FALSE(local_storage->hasBlob(digest2));
    ASSERT_TRUE(remote_storage->hasBlob(digest2));

    std::set<std::string> request_hashes = {digest1.hash(), digest2.hash()};

    BatchReadBlobsRequest request;
    request.set_instance_name(instance_name);
    request.add_digests()->CopyFrom(digest1);
    request.add_digests()->CopyFrom(digest2);
    ASSERT_EQ(request.digests_size(), 2);

    BatchReadBlobsResponse response;
    cas_servicer->BatchReadBlobs(&server_context, &request, &response);
    ASSERT_EQ(response.responses().size(), 2);

    // the local read count is two as a read is triggered in `remote_storage`
    // for the blob that is not in `local_storage`, as well as the one for the
    // blob that is in `local_storage`.

    // Only digest2 is downloaded remotely, while digest1 gets read from the
    // local-cache and digest2 from the remotes local-cache
    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_BATCH_READ,
         buildboxcommonmetrics::CountingMetricValue(2 + 1)},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest2.size_bytes())},
        {MetricNames::COUNTER_NAME_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest1.size_bytes() +
                                                    digest2.size_bytes())},
        {MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
         buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL,
         buildboxcommon::buildboxcommonmetrics::CountingMetricValue(2)},
    };

    validateMetrics(expected);

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::allCollectedByNameWithValues<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>({
            {MetricNames::DISTRIBUTION_NAME_CAS_BATCH_READ_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 digest1.size_bytes())},
            {MetricNames::DISTRIBUTION_NAME_CAS_BATCH_READ_BLOBS_SIZES,
             buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                 digest2.size_bytes())},
        }));

    for (const auto &r : response.responses()) {
        ASSERT_EQ(r.status().code(), grpc::StatusCode::OK);
        // The hash is what we asked for:
        ASSERT_EQ(request_hashes.count(r.digest().hash()), 1);
        // The data corresponds to that blob in its corresponding storage:
        if (r.digest().hash() == digest1.hash()) {
            ASSERT_EQ(r.data(), readBlobFromLocalStorage(r.digest()));
        }
        else {
            ASSERT_EQ(r.data(), readBlobFromRemoteStorage(r.digest()));
        }
    }

    ASSERT_TRUE(local_storage->hasBlob(digest2));
}

TEST_P(CasProxyFixture, BatchReadBlobsSizedExceeded)
{

    // Trying to fetch a blob that exceeds the gRPC message limit returns an
    // invalid-argument error:
    const std::string big_blob(GRPC_DEFAULT_MAX_RECV_MESSAGE_LENGTH, 'X');
    auto digest = DigestGenerator::hash(big_blob);

    BatchReadBlobsRequest request;
    request.set_instance_name(instance_name);
    request.add_digests()->CopyFrom(digest);

    BatchReadBlobsResponse response;
    const grpc::Status return_status =
        cas_servicer->BatchReadBlobs(&server_context, &request, &response);
    ASSERT_EQ(return_status.error_code(), grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, BatchReadBlobsWithMissing)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    const auto digest1 = DigestGenerator::hash("data1");
    const auto digest2 = DigestGenerator::hash("dataX");

    ASSERT_TRUE(local_storage->hasBlob(digest1));
    ASSERT_FALSE(local_storage->hasBlob(digest2));
    ASSERT_FALSE(remote_storage->hasBlob(digest2));

    BatchReadBlobsRequest request;
    request.set_instance_name(instance_name);
    request.add_digests()->CopyFrom(digest1);
    request.add_digests()->CopyFrom(digest2);

    BatchReadBlobsResponse response;
    cas_servicer->BatchReadBlobs(&server_context, &request, &response);

    // Should get no metric for remote reads, just 1 for the local.
    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue>(
            MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL,
            buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)));

    ASSERT_EQ(response.responses().size(), request.digests_size());

    for (const auto &r : response.responses()) {
        if (r.digest().hash() == digest1.hash()) {
            ASSERT_EQ(r.status().code(), grpc::StatusCode::OK);
            ASSERT_EQ(r.data(), readBlobFromLocalStorage(r.digest()));
        }
        else if (r.digest().hash() == digest2.hash()) {
            ASSERT_EQ(r.status().code(), grpc::StatusCode::NOT_FOUND);
        }
        else {
            FAIL() << "Digest in response was not requested.";
        }
    }
}

TEST_P(CasProxyFixture, GetTreeWithMissingRoot)
{
    Digest root_digest;
    root_digest.set_hash("this-digest-is-not-in-the-CAS");
    ASSERT_FALSE(local_storage->hasBlob(root_digest));
    ASSERT_FALSE(remote_storage->hasBlob(root_digest));

    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_digest);

    auto reader = cas_stub->GetTree(&client_context, request);

    GetTreeResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::NOT_FOUND);
}

void assertResponsesContainWholeTree(
    const std::vector<GetTreeResponse> responses,
    const Digest &subdirectory_digest)
{
    /* Helper that asserts that the union `GetTreeResponse`s in the list
     * contain the whole tree without regard for the ordering.
     */
    bool found_root = false, found_subdir = false;
    for (const GetTreeResponse &response : responses) {
        for (const auto &node : response.directories()) {
            if (node.directories_size() > 0) {
                // This node is `root/`:
                found_root = true;
                ASSERT_EQ(node.directories_size(), 1);
                ASSERT_EQ(node.files_size(), 1);
                ASSERT_EQ(node.files(0).name(), "file1.sh");
                ASSERT_TRUE(node.files(0).is_executable());

                // So `root/subdir1` is there:
                const auto subdir1 = node.directories(0);
                ASSERT_EQ(subdir1.name(), "subdir1");
                ASSERT_EQ(subdir1.digest().hash(), subdirectory_digest.hash());
            }
            else {
                // This node represents `subdir1/`:
                found_subdir = true;
                ASSERT_EQ(node.directories_size(), 0);
                ASSERT_EQ(node.files_size(), 1);
                ASSERT_EQ(node.files(0).name(), "file2.c");
                ASSERT_FALSE(node.files(0).is_executable());
            }
        }
    }
    ASSERT_TRUE(found_root);
    ASSERT_TRUE(found_subdir);
}

TEST_P(CasProxyFixture, GetTree)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    /* We'll test this directory structure:
     *  root/
     *  |-- file1.sh*
     *  |-- subdir1/
     *           |-- file2.c
     */

    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    ASSERT_FALSE(local_storage->hasBlob(root_directory_digest));
    ASSERT_FALSE(local_storage->hasBlob(subdirectory_digest));

    ASSERT_TRUE(remote_storage->hasBlob(root_directory_digest));
    ASSERT_TRUE(remote_storage->hasBlob(subdirectory_digest));

    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);

    auto reader = cas_stub->GetTree(&client_context, request);

    // We received a valid reply:
    GetTreeResponse response;
    ASSERT_TRUE(reader->Read(&response));

    // And it contains what we want:
    const size_t expectedSize = 2;
    ASSERT_EQ(response.directories_size(),
              expectedSize); // `root/` and `root/subdir1`

    assertResponsesContainWholeTree({response}, subdirectory_digest);

    ASSERT_TRUE(local_storage->hasBlob(root_directory_digest));
    ASSERT_TRUE(local_storage->hasBlob(subdirectory_digest));
    ASSERT_TRUE(buildboxcommonmetrics::collectedByNameWithValue<
                buildboxcommonmetrics::CountingMetricValue>(
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
        buildboxcommonmetrics::CountingMetricValue(expectedSize)));
}

TEST_P(CasProxyFixture, GetPartialTree)
{
    /* We'll test this directory structure:
     *   root/
     *  |-- file1.sh*
     *  |-- subdir1/
     *           |-- file2.c
     *
     * But the server won't have `subdir1/` in its CAS.
     *
     * According to the Remote Execution specification, this is not a
     * problem: the server must return the parts of the tree that are present
     * and ignore the rest.
     */

    Digest root_directory_digest;
    prepareTree(&root_directory_digest);

    ASSERT_FALSE(local_storage->hasBlob(root_directory_digest));
    ASSERT_TRUE(remote_storage->hasBlob(root_directory_digest));

    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);

    auto reader = cas_stub->GetTree(&client_context, request);

    // We received a valid reply:
    GetTreeResponse response;
    ASSERT_TRUE(reader->Read(&response));

    // And it contains what we want:
    ASSERT_EQ(response.directories_size(), 1);

    const auto node = response.directories(0);
    ASSERT_EQ(node.directories_size(), 1);
    ASSERT_EQ(node.files_size(), 1);
    ASSERT_EQ(node.files(0).name(), "file1.sh");
    ASSERT_TRUE(node.files(0).is_executable());

    ASSERT_TRUE(local_storage->hasBlob(root_directory_digest));
}

TEST_P(CasProxyFixture, GetTreeWithTwoPages)
{
    /* We'll test this directory structure:
     *  root/
     *  |-- file1.sh*
     *  |-- subdir1/
     *           |-- file2.c
     *
     * forcing the server to split the reply into two `GetTreeResponse`s.
     *
     * The server is going to generate the two pages and return them in the
     * stream.
     */

    // Creating `root/`:
    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    ASSERT_FALSE(local_storage->hasBlob(root_directory_digest));
    ASSERT_FALSE(local_storage->hasBlob(subdirectory_digest));

    ASSERT_TRUE(remote_storage->hasBlob(root_directory_digest));
    ASSERT_TRUE(remote_storage->hasBlob(subdirectory_digest));

    // Requesting the tree:
    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);
    request.set_page_size(1);

    auto reader = cas_stub->GetTree(&client_context, request);

    // We received two responses:
    GetTreeResponse response1, response2;
    ASSERT_TRUE(reader->Read(&response1));
    ASSERT_TRUE(reader->Read(&response2));

    ASSERT_EQ(response1.directories_size(), 1);
    ASSERT_EQ(response2.directories_size(), 1);

    ASSERT_EQ(response1.next_page_token(), "1");
    ASSERT_EQ(response2.next_page_token(), "");

    assertResponsesContainWholeTree({response1, response2},
                                    subdirectory_digest);

    // The tree is now cached locally:
    ASSERT_TRUE(local_storage->hasBlob(root_directory_digest));
    ASSERT_TRUE(local_storage->hasBlob(subdirectory_digest));
}

TEST_P(CasProxyFixture, GetTreeSkippingPages)
{
    /* We'll test this directory structure:
     *  root/
     *  |-- file1.sh*
     *  |-- subdir1/
     *           |-- file2.c
     *
     * forcing the server to split the reply into two `GetTreeResponse`s.
     *
     * The server is going to generate the two pages and return them in the
     * stream.
     */

    // Creating `root/`:
    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    // We'll first request the whole tree and simulate an error fetching the
    // second page:
    GetTreeResponse page1;
    {
        grpc::ClientContext context;
        GetTreeRequest request;
        request.set_instance_name(instance_name);
        request.mutable_root_digest()->CopyFrom(root_directory_digest);
        request.set_page_size(1);

        auto reader = cas_stub->GetTree(&context, request);

        ASSERT_TRUE(reader->Read(&page1));
    }

    ASSERT_EQ(page1.next_page_token(), "1");

    // Retrying the request but skipping page 1, which we already have:
    GetTreeResponse page2;
    {
        grpc::ClientContext context;
        GetTreeRequest request;
        request.set_instance_name(instance_name);
        request.mutable_root_digest()->CopyFrom(root_directory_digest);
        request.set_page_size(1);
        request.set_page_token(page1.next_page_token());

        auto reader = cas_stub->GetTree(&context, request);

        ASSERT_TRUE(reader->Read(&page2));
        ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);
    }

    ASSERT_EQ(page2.next_page_token(), "");

    // We effectively fetched the complete tree:
    assertResponsesContainWholeTree({page1, page2}, subdirectory_digest);
}

TEST_P(CasProxyFixture, GetTreeWithPageTokenZero)
{
    // Creating `root/`:
    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    // Requesting the tree:
    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);
    request.set_page_size(1);
    request.set_page_token("0");

    // While page 0 exists for a tree, the server never returns it as a
    // `next_page_token`, and according to the REAPI spec:
    // "A page token, [...] must be a value received in a previous
    // [GetTreeResponse]".
    // So the following request should return `INVALID_ARGUMENT`.

    auto reader = cas_stub->GetTree(&client_context, request);

    GetTreeResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, GetTreeWithInvalidPageToken)
{
    // Creating `root/`:
    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    // Requesting the tree:
    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);
    request.set_page_size(1);
    request.set_page_token("-1");

    auto reader = cas_stub->GetTree(&client_context, request);

    GetTreeResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, GetTreeWithNonExistentPageToken)
{
    // Creating `root/`:
    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    // Requesting the tree:
    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);
    request.set_page_size(1);
    request.set_page_token("3"); // This page does not exist.

    auto reader = cas_stub->GetTree(&client_context, request);

    GetTreeResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, BytestreamReadWithInvalidResourceName)
{
    ReadRequest request;
    request.set_resource_name("/root/path/to/something");

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, BytestreamRemoteReadWithNegativeReadOffset)
{
    const auto digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest));

    ReadRequest request;
    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_offset(-1);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OUT_OF_RANGE);
}

TEST_P(CasProxyFixture, BytestreamRemoteReadWithReadOffsetExceedingRealSize)
{
    const std::string data = "data1";
    const auto digest = DigestGenerator::hash(data);
    ASSERT_TRUE(local_storage->hasBlob(digest));

    ReadRequest request;
    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_offset(
        static_cast<google::protobuf::int64>(data.size() + 1));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OUT_OF_RANGE);
}

TEST_P(CasProxyFixture, BytestreamLocalReadWithNegativeReadOffset)
{
    const auto digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest));

    ReadRequest request;
    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_offset(-1);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OUT_OF_RANGE);
}

TEST_P(CasProxyFixture, BytestreamLocalReadWithNegativeReadLimit)
{
    const std::string data = "data1";
    const Digest digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_limit(-3);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, BytestreamLocalReadWithReadOffsetExceedingRealSize)
{
    const std::string data = "data1";
    const auto digest = DigestGenerator::hash(data);
    ASSERT_TRUE(local_storage->hasBlob(digest));

    ReadRequest request;
    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_offset(
        static_cast<google::protobuf::int64>(data.size() + 1));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OUT_OF_RANGE);
}
TEST_P(CasProxyFixture, BytestreamLocalReadMissingBlob)
{
    const Digest digest = DigestGenerator::hash("data25");
    ASSERT_FALSE(local_storage->hasBlob(digest));

    ReadRequest request;
    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::NOT_FOUND);
}

TEST_P(CasProxyFixture, BytestreamLocalReadEmptyBlob)
{
    const Digest digest = DigestGenerator::hash("");

    local_storage->writeBlob(digest, "");
    ASSERT_TRUE(local_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyFixture, BytestreamLocalRead)
{
    const std::string data = "data1";
    const Digest digest = DigestGenerator::hash(data);
    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), data);
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NAME_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())}};

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>(
            MetricNames::DISTRIBUTION_NAME_CAS_BYTESTREAM_READ_BLOB_SIZES,
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                digest.size_bytes())));

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamLocalLargeRead)
{
    // This blob will have to be split into multiple chunks due to the gRPC's
    // message size limit:
    // Create an empty data array
    constexpr unsigned long sixteen_megabytes = 16UL * 1024UL * 1024UL;
    std::string data = std::string("");
    data.reserve(sixteen_megabytes);

    // Fill data with a repeating NUL terminated string
    const int somedataLength = 9;
    const std::string somedata = std::string("SomeData\0", somedataLength);
    while (data.size() + somedataLength <= sixteen_megabytes) {
        data += somedata;
    }

    const Digest digest = DigestGenerator::hash(data);
    local_storage->writeBlob(digest, data);
    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    std::string fetched_data;
    while (reader->Read(&response)) {
        fetched_data += response.data();
    }
    ASSERT_EQ(DigestGenerator::hash(fetched_data), digest);

    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyFixture, BytestreamLocalReadWithLimit)
{
    const std::string data = "data1";
    const Digest digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_limit(2);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "da");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NAME_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(2)}};

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamLocalReadWithOffset)
{
    const std::string data = "data1";
    const Digest digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_offset(3);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "a1");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NAME_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(2)}};

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamLocalReadWithOffsetAndLimit)
{
    const std::string data = "data1";
    const Digest digest = DigestGenerator::hash("data1");
    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_limit(2);
    request.set_read_offset(1);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "at");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NAME_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(2)},
    };

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamWriteWithInvalidResourceName)
{
    WriteRequest request;
    request.set_resource_name("/root/path/to/something");
    request.set_data("data");
    request.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request));
    ASSERT_EQ(response.committed_size(), 0);
    ASSERT_EQ(writer->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
}

TEST_P(CasProxyFixture, BytestreamRemoteReadMissingBlob)
{
    const Digest digest = DigestGenerator::hash("data25");
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    ReadRequest request;
    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::NOT_FOUND);
}

TEST_P(CasProxyFixture, BytestreamRemoteRead)
{
    const std::string data = "remotedata3";
    const Digest digest = DigestGenerator::hash(data);
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), data);
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    ASSERT_TRUE(local_storage->hasBlob(digest));

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())}};

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>(
            MetricNames::DISTRIBUTION_NAME_CAS_BYTESTREAM_READ_BLOB_SIZES,
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                digest.size_bytes())));

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamRemoteLargeRead)
{
    // This blob will have to be split into multiple chunks due to the gRPC's
    // message size limit:
    // Create an empty data array
    constexpr unsigned long sixteen_megabytes = 16UL * 1024UL * 1024UL;
    std::string data = std::string("");
    data.reserve(sixteen_megabytes);

    // Fill data with a repeating NUL terminated string
    const int somedataLength = 9;
    const std::string somedata = std::string("SomeData\0", somedataLength);
    while (data.size() + somedataLength <= sixteen_megabytes) {
        data += somedata;
    }

    const Digest digest = DigestGenerator::hash(data);
    remote_storage->writeBlob(digest, data);
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    std::string fetched_data;
    while (reader->Read(&response)) {
        fetched_data += response.data();
    }
    ASSERT_EQ(DigestGenerator::hash(fetched_data), digest);

    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    // the digest is written twice, once in the remote's local-cache to
    // prepopulate it, and then once in the local cache when we read it
    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes() * 2)},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())}};

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamRemoteReadWithLimit)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    const std::string data = "remotedata3";
    const Digest digest = DigestGenerator::hash(data);
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_limit(2);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "re");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    // The whole blob is fetched from the remote and written to the
    // local-cache, even if only a subset is returned to the client
    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL,
         buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)}};

    validateMetrics(expected);

    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_EQ(readBlobFromLocalStorage(digest), data);
}

TEST_P(CasProxyFixture, BytestreamRemoteReadWithOffset)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    const std::string data = "remotedata3";
    const Digest digest = DigestGenerator::hash(data);
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_offset(6);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "data3");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    // The whole blob is fetched from the remote and written to the
    // local-cache, even if only a subset is returned to the client
    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
         buildboxcommon::buildboxcommonmetrics::CountingMetricValue(1)}};

    validateMetrics(expected);

    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_EQ(readBlobFromLocalStorage(digest), data);
}

TEST_P(CasProxyFixture, BytestreamRemoteReadWithOffsetAndLimit)
{
    const std::string data = "remotedata3";
    const Digest digest = DigestGenerator::hash(data);
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    request.set_read_limit(2);
    request.set_read_offset(1);

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), "em");
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);

    // The whole blob is fetched from the remote and written to the
    // local-cache, even if only a subset is returned to the client
    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(1)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())}};

    validateMetrics(expected);
}

TEST_P(CasProxyFixture, BytestreamResourceNameIgnoreValuesAfterSize)
{
    const Digest digest = DigestGenerator::hash("remotedata3");
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()) +
                              "/extra/path/to/file.c");

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_TRUE(reader->Read(&response));
    ASSERT_EQ(response.data(), readBlobFromLocalStorage(digest));
    ASSERT_EQ(reader->Finish().error_code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyFixture, BytestreamWrite)
{
    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();
    const std::string data = "data10";
    const Digest digest = DigestGenerator::hash(data);

    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    WriteRequest request;
    request.set_resource_name(resource_name_prefix +
                              "uploads/uuid-goes-here/blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));
    request.set_data(data);
    request.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request));
    ASSERT_TRUE(writer->WritesDone());
    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);

    // Should have triggered no local read metrics,
    // but a local/remote write
    const std::vector<std::string> missingMetrics = {
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL};

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes() * 2)},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())},
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
         buildboxcommonmetrics::CountingMetricValue(1)}};

    ASSERT_TRUE(
        buildboxcommonmetrics::allCollectedByNameWithValuesAndAllMissingByName<
            buildboxcommonmetrics::CountingMetricValue>(expected,
                                                        missingMetrics));

    ASSERT_TRUE(
        buildboxcommon::buildboxcommonmetrics::collectedByNameWithValue<
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue>(
            MetricNames::DISTRIBUTION_NAME_CAS_BYTESTREAM_WRITE_BLOB_SIZES,
            buildboxcommon::buildboxcommonmetrics::DistributionMetricValue(
                digest.size_bytes())));

    ASSERT_EQ(response.committed_size(),
              readBlobFromLocalStorage(digest).size());

    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ASSERT_EQ(readBlobFromLocalStorage(digest), data);
    ASSERT_EQ(readBlobFromRemoteStorage(digest), data);
}

TEST_P(CasProxyFixture, BytestreamWriteMultiPart)
{
    const std::string data = "Part1|Part2";
    const Digest digest = DigestGenerator::hash(data);

    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());
    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data("Part1|");
    request1.set_finish_write(false);

    WriteRequest request2;
    request2.set_resource_name(resource_name);
    request2.set_data("Part2");
    request2.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);

    MetricVector<buildboxcommonmetrics::CountingMetricValue> expected = {
        {MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
         buildboxcommonmetrics::CountingMetricValue(2)},
        {MetricNames::COUNTER_NAME_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes() * 2)},
        {MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
         buildboxcommonmetrics::CountingMetricValue(digest.size_bytes())}};

    validateMetrics(expected);

    ASSERT_EQ(response.committed_size(),
              local_storage->readBlob(digest)->size());

    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ASSERT_EQ(readBlobFromLocalStorage(digest), data);
    ASSERT_EQ(readBlobFromRemoteStorage(digest), data);
}

TEST_P(CasProxyFixture, BytestreamWriteExistingBlob)
{
    // According to the spec:
    // "When attempting an upload, if another client has already completed the
    // upload (which may occur in the middle of a single upload if another
    // client uploads the same blob concurrently), the request will terminate
    // immediately with a response whose `committed_size` is the full size of
    // the uploaded file (regardless of how much data was transmitted by the
    // client)"

    const std::string data = "SomeBlob1234";
    const Digest digest = DigestGenerator::hash(data);

    local_storage->writeBlob(digest, data);
    ASSERT_TRUE(local_storage->hasBlob(digest));

    const auto resourceName = resource_name_prefix +
                              "uploads/uuid-goes-here/blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes());

    // Sending a request containing only a part of the blob:
    WriteRequest request1;
    request1.set_resource_name(resourceName);
    request1.set_data(data.substr(0, 1));
    request1.set_finish_write(false);

    buildboxcommon::buildboxcommonmetrics::clearAllMetricCollection();

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);
    ASSERT_TRUE(writer->Write(request1));

    // The server reports that the whole blob is successfully uploaded, even
    // though we sent just a portion of it:
    ASSERT_TRUE(writer->WritesDone());
    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);

    ASSERT_EQ(response.committed_size(), digest.size_bytes());

    const std::vector<std::string> missingMetrics = {
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
        MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_LOCAL,
    };

    const MetricVector<buildboxcommonmetrics::CountingMetricValue> expected =
        {};

    ASSERT_TRUE(
        buildboxcommonmetrics::allCollectedByNameWithValuesAndAllMissingByName<
            buildboxcommonmetrics::CountingMetricValue>(expected,
                                                        missingMetrics));

    // Data is still consistent:
    ASSERT_TRUE(local_storage->hasBlob(digest));
    ASSERT_EQ(*local_storage->readBlob(digest), data);
}

TEST_P(CasProxyFixture, BytestreamWriteMultiPartWithoutFinishWriteFails)
{
    const std::string data = "Part1|Part2";
    const Digest digest = DigestGenerator::hash(data);

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());

    ASSERT_FALSE(local_storage->hasBlob(digest));

    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data("Part1|");
    request1.set_finish_write(false);

    WriteRequest request2;
    request2.set_resource_name(resource_name);
    request2.set_data("Part2");
    request2.set_finish_write(false);
    // ^ The last request should set `finish_write`. The write will fail.

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);

    // The write was not committed:
    ASSERT_EQ(response.committed_size(), 0);
    ASSERT_FALSE(local_storage->hasBlob(digest));
}

TEST_P(CasProxyFixture, BytestreamWriteMultiPartInvalidDigestFails)
{
    const std::string data = "Part1|Part2";
    const Digest digest = DigestGenerator::hash(data);

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());
    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data("Part1|");
    request1.set_finish_write(false);

    WriteRequest request2;
    request2.set_resource_name(resource_name);
    request2.set_data("SomeOtherPart");
    request2.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
    ASSERT_EQ(response.committed_size(), 0);
}

TEST_P(CasProxyFixture, BytestreamWriteWithInvalidDigestFails)
{
    const std::string data = "data11";
    Digest digest;
    digest.set_hash("hash");

    const auto invalid_data_size =
        static_cast<google::protobuf::int64>(data.size() + 3);
    digest.set_size_bytes(invalid_data_size);

    WriteRequest request;
    request.set_resource_name(resource_name_prefix +
                              "uploads/uuid-goes-here/blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));
    request.set_data(data);
    request.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
    ASSERT_EQ(response.committed_size(), 0);
}

TEST_P(CasProxyFixture, BytestreamWriteExtraRequestFails)
{
    const std::string data = "data10";
    const Digest digest = DigestGenerator::hash(data);

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());
    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data(data);
    request1.set_finish_write(true);
    // (After setting `finish_write` it is illegal to send more requests.)

    WriteRequest request2 = request1;

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
    ASSERT_EQ(response.committed_size(), 0);
}

TEST_P(CasProxyFixture, BytestreamWriteMismatchedRequests)
{
    const std::string data = "Part1|Part2";
    const Digest digest = DigestGenerator::hash(data);

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());
    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data("Part1|");
    request1.set_finish_write(false);

    WriteRequest request2;
    request2.set_resource_name(resource_name + "1");
    // (This changes the Digest, which must match in all requests.)

    request2.set_data("Part2");
    request2.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(),
              grpc::StatusCode::INVALID_ARGUMENT);
    ASSERT_EQ(response.committed_size(), 0);
}

TEST_P(CasProxyFixture, BytestreamWriteResourceNameOnlyInFirstRequest)
{
    const std::string data = "Part1|Part2";
    const Digest digest = DigestGenerator::hash(data);

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());
    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data("Part1|");
    request1.set_finish_write(false);

    WriteRequest request2;
    // resource name is optional in subsequent requests, leave empty
    request2.set_data("Part2");
    request2.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyFixture, BytestreamWriteResourceNameInAllRequests)
{
    const std::string data = "Part1|Part2";
    const Digest digest = DigestGenerator::hash(data);

    const std::string resource_name =
        resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
        digest.hash() + "/" + std::to_string(digest.size_bytes());
    WriteRequest request1;
    request1.set_resource_name(resource_name);
    request1.set_data("Part1|");
    request1.set_finish_write(false);

    WriteRequest request2;
    request2.set_resource_name(resource_name); // same resource name, 2nd part
    request2.set_data("Part2");
    request2.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request1));
    ASSERT_TRUE(writer->Write(request2));
    ASSERT_TRUE(writer->WritesDone());

    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);
}

/* These tests cover the `FindMissingBlobs()` cache optimization.
 * With it a proxy instance will assume that a blob that is reported as present
 * in a remote will stay that way for a while, which will reduce the number of
 * queries that need to be sent to the remote CAS.
 */

class CasProxyWithFindMissingBlobsCacheFixture : public CasProxyFixture {
  protected:
    CasProxyWithFindMissingBlobsCacheFixture()
    {
        assert(std::get<1>(GetParam()) == FINDMISSINGBLOBS_CACHE_ON);
    }
};

TEST_P(CasProxyWithFindMissingBlobsCacheFixture,
       FindMissingBlobsResponseCached)
{
    const auto remote_digest = DigestGenerator::hash("remotedata4");
    ASSERT_TRUE(remote_storage->hasBlob(remote_digest));

    // Request with a digest that is present in the remote:
    FindMissingBlobsRequest request;
    request.set_instance_name(instance_name);
    const auto r1 = request.add_blob_digests();
    r1->CopyFrom(remote_digest);

    // 1) Issuing an initial `FindMissingBlobs()`, which returns empty:
    {
        FindMissingBlobsResponse response;
        cas_servicer->FindMissingBlobs(&server_context, &request, &response);

        ASSERT_TRUE(response.missing_blob_digests().empty());
    }

    local_storage->deleteBlob(remote_digest);
    remote_storage->deleteBlob(remote_digest);
    // ^ This should cause the blob to immediately start being reported as
    // missing by the remote, but the proxy instance should have stored the
    // previous response in cache, so it will consider it as still present and
    // not report it as missing in its response.

    // 2) Calling `FindMissingBlobs()` again, which still returns empty due to
    // the cache:
    {
        FindMissingBlobsResponse response;
        cas_servicer->FindMissingBlobs(&server_context, &request, &response);

        ASSERT_TRUE(response.missing_blob_digests().empty());
    }
}

TEST_P(CasProxyWithFindMissingBlobsCacheFixture,
       BatchUpdateBlobsUpdatesFindMissingBlobsCache)
{
    const auto data = "SomeBlob";
    const Digest digest = DigestGenerator::hash(data);

    // 1) Uploading a blob:
    {
        BatchUpdateBlobsRequest request;
        request.set_instance_name(instance_name);

        auto entry1 = request.add_requests();
        entry1->mutable_digest()->CopyFrom(digest);
        entry1->set_data(data);

        BatchUpdateBlobsResponse response;
        cas_servicer->BatchUpdateBlobs(&server_context, &request, &response);
        ASSERT_TRUE(local_storage->hasBlob(digest));
    }

    // 2) Now that BatchUpdate() caused the proxy's FindMissingBlobs cache to
    // be updated with its digest, deleting the blob from both proxy and
    // remote:
    local_storage->deleteBlob(digest);
    remote_storage->deleteBlob(digest);

    // 3) Issuing a `FindMissingBlobs()` call. The proxy will check its cache,
    // and assume that the blob will still be present in the remote; skipping
    // forwarding that request.
    {
        // Request with a digest that is present in the remote:
        FindMissingBlobsRequest request;
        request.set_instance_name(instance_name);
        const auto r1 = request.add_blob_digests();
        r1->CopyFrom(digest);

        FindMissingBlobsResponse response;
        cas_servicer->FindMissingBlobs(&server_context, &request, &response);

        ASSERT_TRUE(response.missing_blob_digests().empty());
    }
}

TEST_P(CasProxyWithFindMissingBlobsCacheFixture,
       ByteStreamWriteUpdatesFindMissingBlobsCache)
{
    const auto data = "SomeBlob";
    const Digest digest = DigestGenerator::hash(data);

    // 1) Uploading a blob:
    {
        WriteRequest request;
        request.set_resource_name(
            resource_name_prefix + "uploads/uuid-goes-here/blobs/" +
            digest.hash() + "/" + std::to_string(digest.size_bytes()));
        request.set_data(data);
        request.set_finish_write(true);

        WriteResponse response;
        auto writer = cas_bytestream_stub->Write(&client_context, &response);
        ASSERT_TRUE(writer->Write(request));
        ASSERT_TRUE(writer->WritesDone());
        ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);
    }

    // 2) Now that BatchUpdate() caused the proxy's FindMissingBlobs cache to
    // be updated with its digest, deleting the blob from both proxy and
    // remote:
    local_storage->deleteBlob(digest);
    remote_storage->deleteBlob(digest);

    // 3) Issuing a `FindMissingBlobs()` call. The proxy will check its cache,
    // and assume that the blob will still be present in the remote; not
    // forwarding that request.
    {
        // Request with a digest that is present in the remote:
        FindMissingBlobsRequest request;
        request.set_instance_name(instance_name);
        const auto r1 = request.add_blob_digests();
        r1->CopyFrom(digest);

        FindMissingBlobsResponse response;
        cas_servicer->FindMissingBlobs(&server_context, &request, &response);

        ASSERT_TRUE(response.missing_blob_digests().empty());
    }
}

class CasProxyUnReachableRemoteFixture : public CasProxyFixture {
  protected:
    CasProxyUnReachableRemoteFixture() { stopRemoteServer(); }
};

TEST_P(CasProxyUnReachableRemoteFixture, FindMissingBlobsWithUnreachableRemote)
{
    FindMissingBlobsRequest request;
    request.set_instance_name(instance_name);

    const auto r = request.add_blob_digests();
    r->CopyFrom(DigestGenerator::hash("data1"));

    FindMissingBlobsResponse response;
    grpc::Status status;
    ASSERT_NO_THROW(status = cas_servicer->FindMissingBlobs(
                        &server_context, &request, &response));

    ASSERT_NE(status.error_code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyUnReachableRemoteFixture, GetTreeWithUnreachableRemote)
{
    Digest root_directory_digest, subdirectory_digest;
    prepareTree(&root_directory_digest, &subdirectory_digest);

    GetTreeRequest request;
    request.set_instance_name(instance_name);
    request.mutable_root_digest()->CopyFrom(root_directory_digest);

    auto reader = cas_stub->GetTree(&client_context, request);

    GetTreeResponse response;
    ASSERT_FALSE(reader->Read(&response));
}

TEST_P(CasProxyUnReachableRemoteFixture,
       BytestreamRemoteReadWithUnreachableRemote)
{
    const std::string data = "remotedata3";
    const Digest digest = DigestGenerator::hash(data);
    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_TRUE(remote_storage->hasBlob(digest));

    ReadRequest request;

    request.set_resource_name(resource_name_prefix + "blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));

    auto reader = cas_bytestream_stub->Read(&client_context, request);

    ReadResponse response;
    ASSERT_FALSE(reader->Read(&response));
    ASSERT_NE(reader->Finish().error_code(), grpc::StatusCode::OK);
}

TEST_P(CasProxyUnReachableRemoteFixture, BytestreamWriteWithUnreachableRemote)
{
    const std::string data = "data10";
    const Digest digest = DigestGenerator::hash(data);

    ASSERT_FALSE(local_storage->hasBlob(digest));
    ASSERT_FALSE(remote_storage->hasBlob(digest));

    WriteRequest request;
    request.set_resource_name(resource_name_prefix +
                              "uploads/uuid-goes-here/blobs/" + digest.hash() +
                              "/" + std::to_string(digest.size_bytes()));
    request.set_data(data);
    request.set_finish_write(true);

    WriteResponse response;
    auto writer = cas_bytestream_stub->Write(&client_context, &response);

    ASSERT_TRUE(writer->Write(request));
    ASSERT_TRUE(writer->WritesDone());
    ASSERT_EQ(writer->Finish().error_code(), grpc::StatusCode::OK);
}

#endif // INCLUDED_BUILDBOXCASD_CASPROXY_H
