
#include "build/bazel/remote/execution/v2/remote_execution.pb.h"
#include <gmock/gmock.h>
#include <grpcpp/client_context.h>
#include <gtest/gtest.h>

#include <buildboxcasd_requestcontextmanager.h>
#include <buildboxcommon_protos.h>
#include <buildboxcommon_requestmetadata.h>
#include <map>
#include <string>

using namespace buildboxcasd;
using namespace build::bazel::remote::execution::v2;
using namespace testing;

// Piggy-back on grpc::ClientContext's own testing apparatus to allow
// inspecting the metadata we've added in the test
class grpc::testing::ClientContextTestPeer {
  public:
    explicit ClientContextTestPeer(const grpc::ClientContext &ctx) : ctx_(&ctx)
    {
    }

    std::multimap<std::string, std::string> getMetadata()
    {
        return ctx_->send_initial_metadata_;
    }

  private:
    const grpc::ClientContext *ctx_;
};

TEST(RequestContextManager, GetSetRequestMetadata)
{
    RequestMetadata metadata = RequestMetadata::default_instance();
    metadata.mutable_tool_details()->set_tool_name("test-suite");

    RequestContextManager::setRequestMetadata(metadata);
    ASSERT_EQ(
        RequestContextManager::getRequestMetadata().tool_details().tool_name(),
        "test-suite (via buildbox-casd)");
    RequestContextManager::reset();
}

TEST(RequestContextManager, AttachMetadata)
{
    RequestMetadata metadata = RequestMetadata::default_instance();
    metadata.mutable_tool_details()->set_tool_name("test-suite");

    RequestContextManager::setRequestMetadata(metadata);

    grpc::ClientContext ctx;
    RequestContextManager::attachMetadata(&ctx);
    grpc::testing::ClientContextTestPeer contextPeer(ctx);
    auto attachedMetadata = contextPeer.getMetadata();
    ASSERT_EQ(attachedMetadata.count(
                  buildboxcommon::RequestMetadataGenerator::HEADER_NAME),
              1);

    RequestMetadata expectedMetadata = RequestMetadata::default_instance();
    expectedMetadata.CopyFrom(metadata);
    expectedMetadata.mutable_tool_details()->set_tool_name(
        metadata.tool_details().tool_name() + " (via buildbox-casd)");
    if (auto header = attachedMetadata.find(
            buildboxcommon::RequestMetadataGenerator::HEADER_NAME);
        header != attachedMetadata.end()) {
        ASSERT_EQ(header->second, expectedMetadata.SerializeAsString());
    }
    RequestContextManager::reset();
}

TEST(RequestContextManager, MissingMetadataNotAttached)
{
    // Calling RequestContextManager::attachMetadata without setting the
    // metadata for this thread should result in no metadata ever being
    // attached.
    grpc::ClientContext ctx;
    grpc::testing::ClientContextTestPeer contextPeer(ctx);
    auto attachedMetadata = contextPeer.getMetadata();
    ASSERT_EQ(attachedMetadata.count(
                  buildboxcommon::RequestMetadataGenerator::HEADER_NAME),
              0);
    RequestContextManager::attachMetadata(&ctx);
    attachedMetadata = contextPeer.getMetadata();
    ASSERT_EQ(attachedMetadata.count(
                  buildboxcommon::RequestMetadataGenerator::HEADER_NAME),
              0);
    RequestContextManager::reset();
}
