unittests/Telemetry/TelemetryTest.cpp

//===- llvm/unittest/Telemetry/TelemetryTest.cpp - Telemetry unittests ---===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "llvm/Telemetry/Telemetry.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Error.h"
#include "gtest/gtest.h"
#include <optional>
#include <vector>

namespace llvm {
namespace telemetry {
// Testing parameters.
//
// These are set by each test to force certain outcomes.
struct TestContext {
  // Controlling whether there is vendor plugin.  In "real" implementation, the
  // plugin-registration framework will handle the overrides but for tests, we
  // just use a bool flag to decide which function to call.
  bool HasVendorPlugin = false;

  // This field contains data emitted by the framework for later
  // verification by the tests.
  std::string Buffer = "";

  // The expected Uuid generated by the fake tool.
  std::string ExpectedUuid = "";
};

class StringSerializer : public Serializer {
public:
  const std::string &getString() { return Buffer; }

  Error init() override {
    if (Started)
      return createStringError("Serializer already in use");
    Started = true;
    Buffer.clear();
    return Error::success();
  }

  void write(StringRef KeyName, bool Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, StringRef Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, int Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, long Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, long long Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, unsigned int Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, unsigned long Value) override {
    writeHelper(KeyName, Value);
  }

  void write(StringRef KeyName, unsigned long long Value) override {
    writeHelper(KeyName, Value);
  }

  void beginObject(StringRef KeyName) override {
    Children.push_back(std::string("\n"));
    ChildrenNames.push_back(KeyName.str());
  }

  void endObject() override {
    assert(!Children.empty() && !ChildrenNames.empty());
    std::string ChildBuff = Children.back();
    std::string Name = ChildrenNames.back();
    Children.pop_back();
    ChildrenNames.pop_back();
    writeHelper(Name, ChildBuff);
  }

  Error finalize() override {
    assert(Children.empty() && ChildrenNames.empty());
    if (!Started)
      return createStringError("Serializer not currently in use");
    Started = false;
    return Error::success();
  }

private:
  template <typename T> void writeHelper(StringRef Name, T Value) {
    assert(Started && "serializer not started");
    if (Children.empty())
      Buffer.append((Name + ":" + Twine(Value) + "\n").str());
    else
      Children.back().append((Name + ":" + Twine(Value) + "\n").str());
  }

  bool Started = false;
  std::string Buffer;
  std::vector<std::string> Children;
  std::vector<std::string> ChildrenNames;
};

namespace vendor {
struct VendorConfig : public Config {
  VendorConfig(bool Enable) : Config(Enable) {}
  std::optional<std::string> makeSessionId() override {
    static int seed = 0;
    return std::to_string(seed++);
  }
};

std::shared_ptr<Config> getTelemetryConfig(const TestContext &Ctxt) {
  return std::make_shared<VendorConfig>(/*EnableTelemetry=*/true);
}

class TestStorageDestination : public Destination {
public:
  TestStorageDestination(TestContext *Ctxt) : CurrentContext(Ctxt) {}

  Error receiveEntry(const TelemetryInfo *Entry) override {
    if (Error Err = serializer.init())
      return Err;

    Entry->serialize(serializer);
    if (Error Err = serializer.finalize())
      return Err;

    CurrentContext->Buffer.append(serializer.getString());
    return Error::success();
  }

  StringLiteral name() const override { return "TestDestination"; }

private:
  TestContext *CurrentContext;
  StringSerializer serializer;
};

struct StartupInfo : public TelemetryInfo {
  std::string ToolName;
  std::map<std::string, std::string> MetaData;

  void serialize(Serializer &serializer) const override {
    TelemetryInfo::serialize(serializer);
    serializer.write("ToolName", ToolName);
    serializer.write("MetaData", MetaData);
  }
};

struct ExitInfo : public TelemetryInfo {
  int ExitCode;
  std::string ExitDesc;
  void serialize(Serializer &serializer) const override {
    TelemetryInfo::serialize(serializer);
    serializer.write("ExitCode", ExitCode);
    serializer.write("ExitDesc", ExitDesc);
  }
};

class TestManager : public Manager {
public:
  static std::unique_ptr<TestManager>
  createInstance(Config *Config, TestContext *CurrentContext) {
    if (!Config->EnableTelemetry)
      return nullptr;
    CurrentContext->ExpectedUuid = *(Config->makeSessionId());
    std::unique_ptr<TestManager> Ret = std::make_unique<TestManager>(
        CurrentContext, CurrentContext->ExpectedUuid);

    // Add a destination.
    Ret->addDestination(
        std::make_unique<TestStorageDestination>(CurrentContext));

    return Ret;
  }

  TestManager(TestContext *Ctxt, std::string Id)
      : CurrentContext(Ctxt), SessionId(Id) {}

  Error preDispatch(TelemetryInfo *Entry) override {
    Entry->SessionId = SessionId;
    (void)CurrentContext;
    return Error::success();
  }

  std::string getSessionId() { return SessionId; }

private:
  TestContext *CurrentContext;
  const std::string SessionId;
};
} // namespace vendor

std::shared_ptr<Config> getTelemetryConfig(const TestContext &Ctxt) {
  if (Ctxt.HasVendorPlugin)
    return vendor::getTelemetryConfig(Ctxt);

  return std::make_shared<Config>(false);
}

TEST(TelemetryTest, TelemetryDisabled) {
  TestContext Context;
  Context.HasVendorPlugin = false;

  std::shared_ptr<Config> Config = getTelemetryConfig(Context);
  auto Manager = vendor::TestManager::createInstance(Config.get(), &Context);
  EXPECT_EQ(nullptr, Manager);
}

TEST(TelemetryTest, TelemetryEnabled) {
  const std::string ToolName = "TelemetryTestTool";

  // Preset some params.
  TestContext Context;
  Context.HasVendorPlugin = true;
  Context.Buffer.clear();

  std::shared_ptr<Config> Config = getTelemetryConfig(Context);
  auto Manager = vendor::TestManager::createInstance(Config.get(), &Context);

  EXPECT_STREQ(Manager->getSessionId().c_str(), Context.ExpectedUuid.c_str());

  vendor::StartupInfo S;
  S.ToolName = ToolName;
  S.MetaData["a"] = "A";
  S.MetaData["b"] = "B";

  Error startupEmitStatus = Manager->dispatch(&S);
  EXPECT_FALSE(startupEmitStatus);
  std::string ExpectedBuffer =
      "SessionId:0\nToolName:TelemetryTestTool\nMetaData:\na:A\nb:B\n\n";
  EXPECT_EQ(ExpectedBuffer, Context.Buffer);
  Context.Buffer.clear();

  vendor::ExitInfo E;
  E.ExitCode = 0;
  E.ExitDesc = "success";
  Error exitEmitStatus = Manager->dispatch(&E);
  EXPECT_FALSE(exitEmitStatus);
  ExpectedBuffer = "SessionId:0\nExitCode:0\nExitDesc:success\n";
  EXPECT_EQ(ExpectedBuffer, Context.Buffer);
}

} // namespace telemetry
} // namespace llvm