266 lines · cpp
1//===- llvm/unittest/Telemetry/TelemetryTest.cpp - Telemetry unittests ---===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/Telemetry/Telemetry.h"10#include "llvm/ADT/StringRef.h"11#include "llvm/Support/Casting.h"12#include "llvm/Support/Error.h"13#include "gtest/gtest.h"14#include <optional>15#include <vector>16 17namespace llvm {18namespace telemetry {19// Testing parameters.20//21// These are set by each test to force certain outcomes.22struct TestContext {23 // Controlling whether there is vendor plugin. In "real" implementation, the24 // plugin-registration framework will handle the overrides but for tests, we25 // just use a bool flag to decide which function to call.26 bool HasVendorPlugin = false;27 28 // This field contains data emitted by the framework for later29 // verification by the tests.30 std::string Buffer = "";31 32 // The expected Uuid generated by the fake tool.33 std::string ExpectedUuid = "";34};35 36class StringSerializer : public Serializer {37public:38 const std::string &getString() { return Buffer; }39 40 Error init() override {41 if (Started)42 return createStringError("Serializer already in use");43 Started = true;44 Buffer.clear();45 return Error::success();46 }47 48 void write(StringRef KeyName, bool Value) override {49 writeHelper(KeyName, Value);50 }51 52 void write(StringRef KeyName, StringRef Value) override {53 writeHelper(KeyName, Value);54 }55 56 void write(StringRef KeyName, int Value) override {57 writeHelper(KeyName, Value);58 }59 60 void write(StringRef KeyName, long Value) override {61 writeHelper(KeyName, Value);62 }63 64 void write(StringRef KeyName, long long Value) override {65 writeHelper(KeyName, Value);66 }67 68 void write(StringRef KeyName, unsigned int Value) override {69 writeHelper(KeyName, Value);70 }71 72 void write(StringRef KeyName, unsigned long Value) override {73 writeHelper(KeyName, Value);74 }75 76 void write(StringRef KeyName, unsigned long long Value) override {77 writeHelper(KeyName, Value);78 }79 80 void beginObject(StringRef KeyName) override {81 Children.push_back(std::string("\n"));82 ChildrenNames.push_back(KeyName.str());83 }84 85 void endObject() override {86 assert(!Children.empty() && !ChildrenNames.empty());87 std::string ChildBuff = Children.back();88 std::string Name = ChildrenNames.back();89 Children.pop_back();90 ChildrenNames.pop_back();91 writeHelper(Name, ChildBuff);92 }93 94 Error finalize() override {95 assert(Children.empty() && ChildrenNames.empty());96 if (!Started)97 return createStringError("Serializer not currently in use");98 Started = false;99 return Error::success();100 }101 102private:103 template <typename T> void writeHelper(StringRef Name, T Value) {104 assert(Started && "serializer not started");105 if (Children.empty())106 Buffer.append((Name + ":" + Twine(Value) + "\n").str());107 else108 Children.back().append((Name + ":" + Twine(Value) + "\n").str());109 }110 111 bool Started = false;112 std::string Buffer;113 std::vector<std::string> Children;114 std::vector<std::string> ChildrenNames;115};116 117namespace vendor {118struct VendorConfig : public Config {119 VendorConfig(bool Enable) : Config(Enable) {}120 std::optional<std::string> makeSessionId() override {121 static int seed = 0;122 return std::to_string(seed++);123 }124};125 126std::shared_ptr<Config> getTelemetryConfig(const TestContext &Ctxt) {127 return std::make_shared<VendorConfig>(/*EnableTelemetry=*/true);128}129 130class TestStorageDestination : public Destination {131public:132 TestStorageDestination(TestContext *Ctxt) : CurrentContext(Ctxt) {}133 134 Error receiveEntry(const TelemetryInfo *Entry) override {135 if (Error Err = serializer.init())136 return Err;137 138 Entry->serialize(serializer);139 if (Error Err = serializer.finalize())140 return Err;141 142 CurrentContext->Buffer.append(serializer.getString());143 return Error::success();144 }145 146 StringLiteral name() const override { return "TestDestination"; }147 148private:149 TestContext *CurrentContext;150 StringSerializer serializer;151};152 153struct StartupInfo : public TelemetryInfo {154 std::string ToolName;155 std::map<std::string, std::string> MetaData;156 157 void serialize(Serializer &serializer) const override {158 TelemetryInfo::serialize(serializer);159 serializer.write("ToolName", ToolName);160 serializer.write("MetaData", MetaData);161 }162};163 164struct ExitInfo : public TelemetryInfo {165 int ExitCode;166 std::string ExitDesc;167 void serialize(Serializer &serializer) const override {168 TelemetryInfo::serialize(serializer);169 serializer.write("ExitCode", ExitCode);170 serializer.write("ExitDesc", ExitDesc);171 }172};173 174class TestManager : public Manager {175public:176 static std::unique_ptr<TestManager>177 createInstance(Config *Config, TestContext *CurrentContext) {178 if (!Config->EnableTelemetry)179 return nullptr;180 CurrentContext->ExpectedUuid = *(Config->makeSessionId());181 std::unique_ptr<TestManager> Ret = std::make_unique<TestManager>(182 CurrentContext, CurrentContext->ExpectedUuid);183 184 // Add a destination.185 Ret->addDestination(186 std::make_unique<TestStorageDestination>(CurrentContext));187 188 return Ret;189 }190 191 TestManager(TestContext *Ctxt, std::string Id)192 : CurrentContext(Ctxt), SessionId(Id) {}193 194 Error preDispatch(TelemetryInfo *Entry) override {195 Entry->SessionId = SessionId;196 (void)CurrentContext;197 return Error::success();198 }199 200 std::string getSessionId() { return SessionId; }201 202private:203 TestContext *CurrentContext;204 const std::string SessionId;205};206} // namespace vendor207 208std::shared_ptr<Config> getTelemetryConfig(const TestContext &Ctxt) {209 if (Ctxt.HasVendorPlugin)210 return vendor::getTelemetryConfig(Ctxt);211 212 return std::make_shared<Config>(false);213}214 215#if LLVM_ENABLE_TELEMETRY216#define TELEMETRY_TEST(suite, test) TEST(suite, test)217#else218#define TELEMETRY_TEST(suite, test) TEST(DISABLED_##suite, test)219#endif220 221TELEMETRY_TEST(TelemetryTest, TelemetryDisabled) {222 TestContext Context;223 Context.HasVendorPlugin = false;224 225 std::shared_ptr<Config> Config = getTelemetryConfig(Context);226 auto Manager = vendor::TestManager::createInstance(Config.get(), &Context);227 EXPECT_EQ(nullptr, Manager);228}229 230TELEMETRY_TEST(TelemetryTest, TelemetryEnabled) {231 const std::string ToolName = "TelemetryTestTool";232 233 // Preset some params.234 TestContext Context;235 Context.HasVendorPlugin = true;236 Context.Buffer.clear();237 238 std::shared_ptr<Config> Config = getTelemetryConfig(Context);239 auto Manager = vendor::TestManager::createInstance(Config.get(), &Context);240 241 EXPECT_STREQ(Manager->getSessionId().c_str(), Context.ExpectedUuid.c_str());242 243 vendor::StartupInfo S;244 S.ToolName = ToolName;245 S.MetaData["a"] = "A";246 S.MetaData["b"] = "B";247 248 Error startupEmitStatus = Manager->dispatch(&S);249 EXPECT_FALSE(startupEmitStatus);250 std::string ExpectedBuffer =251 "SessionId:0\nToolName:TelemetryTestTool\nMetaData:\na:A\nb:B\n\n";252 EXPECT_EQ(ExpectedBuffer, Context.Buffer);253 Context.Buffer.clear();254 255 vendor::ExitInfo E;256 E.ExitCode = 0;257 E.ExitDesc = "success";258 Error exitEmitStatus = Manager->dispatch(&E);259 EXPECT_FALSE(exitEmitStatus);260 ExpectedBuffer = "SessionId:0\nExitCode:0\nExitDesc:success\n";261 EXPECT_EQ(ExpectedBuffer, Context.Buffer);262}263 264} // namespace telemetry265} // namespace llvm266