Make MockTracer's behavior more customizable. (#59)

rnburn · web-flow · commit f3c1f42601d1 · 2018-02-27T15:23:36.000-05:00
diff --git a/mocktracer/include/opentracing/mocktracer/tracer.h b/mocktracer/include/opentracing/mocktracer/tracer.h
@@ -11,17 +11,35 @@ namespace opentracing {
 BEGIN_OPENTRACING_ABI_NAMESPACE
 namespace mocktracer {
 
+struct PropagationOptions {
+  // Specifies what key to use when injecting and extracting span context.
+  std::string propagation_key = "x-ot_span_context";
+
+  // If inject_error_code is non-zero, MockTracer::Inject fails with
+  // inject_error_code.
+  std::error_code inject_error_code;
+
+  // If extract_error_code is non-zero, MockTracer::Extract fails with
+  // extract_error_code.
+  std::error_code extract_error_code;
+};
+
 struct MockTracerOptions {
+  // Recorder is sent spans when they are finished. If nullptr, all finished
+  // spans are dropped.
   std::unique_ptr<Recorder> recorder;
+
+  // PropagationOptions allows you to customize how the mocktracer's SpanContext
+  // is propagated.
+  PropagationOptions propagation_options;
 };
 
 // MockTracer provides implements the OpenTracing Tracer API. It provides
 // convenient access to finished spans in such a way as to support testing.
 class MockTracer : public Tracer,
                    public std::enable_shared_from_this<MockTracer> {
  public:
-  explicit MockTracer(MockTracerOptions&& options)
-      : recorder_{std::move(options.recorder)} {}
+  explicit MockTracer(MockTracerOptions&& options);
 
   std::unique_ptr<Span> StartSpanWithOptions(
       string_view operation_name, const StartSpanOptions& options) const
@@ -54,6 +72,7 @@ class MockTracer : public Tracer,
 
  private:
   std::unique_ptr<Recorder> recorder_;
+  PropagationOptions propagation_options_;
   std::mutex mutex_;
   std::vector<SpanData> spans_;
 };
diff --git a/mocktracer/src/mock_span_context.h b/mocktracer/src/mock_span_context.h
@@ -36,15 +36,17 @@ class MockSpanContext : public SpanContext {
   void SetData(SpanContextData& data);
 
   template <class Carrier>
-  expected<void> Inject(Carrier& writer) const {
+  expected<void> Inject(const PropagationOptions& propagation_options,
+                        Carrier& writer) const {
     std::lock_guard<std::mutex> lock_guard{baggage_mutex_};
-    return InjectSpanContext(writer, data_);
+    return InjectSpanContext(propagation_options, writer, data_);
   }
 
   template <class Carrier>
-  expected<bool> Extract(Carrier& reader) {
+  expected<bool> Extract(const PropagationOptions& propagation_options,
+                         Carrier& reader) {
     std::lock_guard<std::mutex> lock_guard{baggage_mutex_};
-    return ExtractSpanContext(reader, data_);
+    return ExtractSpanContext(propagation_options, reader, data_);
   }
 
  private:
diff --git a/mocktracer/src/propagation.cpp b/mocktracer/src/propagation.cpp
@@ -9,8 +9,6 @@
 namespace opentracing {
 BEGIN_OPENTRACING_ABI_NAMESPACE
 namespace mocktracer {
-const opentracing::string_view propagation_key = "x-ot-span-context";
-
 static void WriteString(std::ostream& ostream, const std::string& s) {
   const uint32_t size = static_cast<uint32_t>(s.size());
   ostream.write(reinterpret_cast<const char*>(&size), sizeof(size));
@@ -24,8 +22,9 @@ static void ReadString(std::istream& istream, std::string& s) {
   istream.read(&s[0], size);
 }
 
-expected<void> InjectSpanContext(std::ostream& carrier,
-                                 const SpanContextData& span_context_data) {
+expected<void> InjectSpanContext(
+    const PropagationOptions& /*propagation_options*/, std::ostream& carrier,
+    const SpanContextData& span_context_data) {
   carrier.write(reinterpret_cast<const char*>(&span_context_data.trace_id),
                 sizeof(&span_context_data.trace_id));
   carrier.write(reinterpret_cast<const char*>(&span_context_data.span_id),
@@ -50,8 +49,9 @@ expected<void> InjectSpanContext(std::ostream& carrier,
   return {};
 }
 
-expected<bool> ExtractSpanContext(std::istream& carrier,
-                                  SpanContextData& span_context_data) try {
+expected<bool> ExtractSpanContext(
+    const PropagationOptions& /*propagation_options*/, std::istream& carrier,
+    SpanContextData& span_context_data) try {
   // istream::peek returns EOF if it's in an error state, so check for an error
   // state first before checking for an empty stream.
   if (!carrier.good()) {
@@ -86,10 +86,12 @@ expected<bool> ExtractSpanContext(std::istream& carrier,
       std::make_error_code(std::errc::not_enough_memory));
 }
 
-expected<void> InjectSpanContext(const TextMapWriter& carrier,
+expected<void> InjectSpanContext(const PropagationOptions& propagation_options,
+                                 const TextMapWriter& carrier,
                                  const SpanContextData& span_context_data) {
   std::ostringstream ostream;
-  auto result = InjectSpanContext(ostream, span_context_data);
+  auto result =
+      InjectSpanContext(propagation_options, ostream, span_context_data);
   if (!result) {
     return result;
   }
@@ -103,7 +105,7 @@ expected<void> InjectSpanContext(const TextMapWriter& carrier,
         std::make_error_code(std::errc::not_enough_memory));
   }
 
-  result = carrier.Set(propagation_key, context_value);
+  result = carrier.Set(propagation_options.propagation_key, context_value);
   if (!result) {
     return result;
   }
@@ -142,9 +144,11 @@ static opentracing::expected<opentracing::string_view> LookupKey(
 
 template <class KeyCompare>
 static opentracing::expected<bool> ExtractSpanContext(
+    const PropagationOptions& propagation_options,
     const opentracing::TextMapReader& carrier,
     SpanContextData& span_context_data, KeyCompare key_compare) {
-  auto value_maybe = LookupKey(carrier, propagation_key, key_compare);
+  auto value_maybe =
+      LookupKey(carrier, propagation_options.propagation_key, key_compare);
   if (!value_maybe) {
     if (value_maybe.error() == opentracing::key_not_found_error) {
       return false;
@@ -165,22 +169,26 @@ static opentracing::expected<bool> ExtractSpanContext(
         opentracing::span_context_corrupted_error);
   }
   std::istringstream istream{base64_decoding};
-  return ExtractSpanContext(istream, span_context_data);
+  return ExtractSpanContext(propagation_options, istream, span_context_data);
 }
 
-expected<bool> ExtractSpanContext(const TextMapReader& carrier,
+expected<bool> ExtractSpanContext(const PropagationOptions& propagation_options,
+                                  const TextMapReader& carrier,
                                   SpanContextData& span_context_data) {
-  return ExtractSpanContext(carrier, span_context_data,
+  return ExtractSpanContext(propagation_options, carrier, span_context_data,
                             std::equal_to<string_view>{});
 }
 
-expected<void> InjectSpanContext(const HTTPHeadersWriter& carrier,
+expected<void> InjectSpanContext(const PropagationOptions& propagation_options,
+                                 const HTTPHeadersWriter& carrier,
                                  const SpanContextData& span_context_data) {
-  return InjectSpanContext(static_cast<const TextMapWriter&>(carrier),
+  return InjectSpanContext(propagation_options,
+                           static_cast<const TextMapWriter&>(carrier),
                            span_context_data);
 }
 
-expected<bool> ExtractSpanContext(const HTTPHeadersReader& carrier,
+expected<bool> ExtractSpanContext(const PropagationOptions& propagation_options,
+                                  const HTTPHeadersReader& carrier,
                                   SpanContextData& span_context_data) {
   auto iequals = [](opentracing::string_view lhs,
                     opentracing::string_view rhs) {
@@ -190,7 +198,8 @@ expected<bool> ExtractSpanContext(const HTTPHeadersReader& carrier,
                         return std::tolower(a) == std::tolower(b);
                       });
   };
-  return ExtractSpanContext(carrier, span_context_data, iequals);
+  return ExtractSpanContext(propagation_options, carrier, span_context_data,
+                            iequals);
 }
 }  // namespace mocktracer
 END_OPENTRACING_ABI_NAMESPACE
diff --git a/mocktracer/src/propagation.h b/mocktracer/src/propagation.h
@@ -2,27 +2,35 @@
 #define OPENTRACING_MOCKTRACER_PROPAGATION_H
 
 #include <opentracing/mocktracer/recorder.h>
+#include <opentracing/mocktracer/tracer.h>
 #include <opentracing/propagation.h>
 
 namespace opentracing {
 BEGIN_OPENTRACING_ABI_NAMESPACE
 namespace mocktracer {
-expected<void> InjectSpanContext(std::ostream& carrier,
+
+expected<void> InjectSpanContext(const PropagationOptions& propagation_options,
+                                 std::ostream& carrier,
                                  const SpanContextData& span_context_data);
 
-expected<bool> ExtractSpanContext(std::istream& carrier,
+expected<bool> ExtractSpanContext(const PropagationOptions& propagation_options,
+                                  std::istream& carrier,
                                   SpanContextData& span_context_data);
 
-expected<void> InjectSpanContext(const TextMapWriter& carrier,
+expected<void> InjectSpanContext(const PropagationOptions& propagation_options,
+                                 const TextMapWriter& carrier,
                                  const SpanContextData& span_context_data);
 
-expected<bool> ExtractSpanContext(const TextMapReader& carrier,
+expected<bool> ExtractSpanContext(const PropagationOptions& propagation_options,
+                                  const TextMapReader& carrier,
                                   SpanContextData& span_context_data);
 
-expected<void> InjectSpanContext(const HTTPHeadersWriter& carrier,
+expected<void> InjectSpanContext(const PropagationOptions& propagation_options,
+                                 const HTTPHeadersWriter& carrier,
                                  const SpanContextData& span_context_data);
 
-expected<bool> ExtractSpanContext(const HTTPHeadersReader& carrier,
+expected<bool> ExtractSpanContext(const PropagationOptions& propagation_options,
+                                  const HTTPHeadersReader& carrier,
                                   SpanContextData& span_context_data);
 }  // namespace mocktracer
 END_OPENTRACING_ABI_NAMESPACE
diff --git a/mocktracer/src/tracer.cpp b/mocktracer/src/tracer.cpp
@@ -4,25 +4,33 @@
 
 #include "mock_span.h"
 #include "mock_span_context.h"
+#include "propagation.h"
 
 namespace opentracing {
 BEGIN_OPENTRACING_ABI_NAMESPACE
 namespace mocktracer {
 
 template <class Carrier>
-static expected<void> InjectImpl(const opentracing::SpanContext& span_context,
+static expected<void> InjectImpl(const PropagationOptions& propagation_options,
+                                 const opentracing::SpanContext& span_context,
                                  Carrier& writer) {
+  if (propagation_options.inject_error_code.value() != 0) {
+    return opentracing::make_unexpected(propagation_options.inject_error_code);
+  }
   auto mock_span_context = dynamic_cast<const MockSpanContext*>(&span_context);
   if (mock_span_context == nullptr) {
     return opentracing::make_unexpected(
         opentracing::invalid_span_context_error);
   }
-  return mock_span_context->Inject(writer);
+  return mock_span_context->Inject(propagation_options, writer);
 }
 
 template <class Carrier>
 opentracing::expected<std::unique_ptr<opentracing::SpanContext>> ExtractImpl(
-    Carrier& reader) {
+    const PropagationOptions& propagation_options, Carrier& reader) {
+  if (propagation_options.extract_error_code.value() != 0) {
+    return opentracing::make_unexpected(propagation_options.extract_error_code);
+  }
   MockSpanContext* mock_span_context;
   try {
     mock_span_context = new MockSpanContext{};
@@ -31,7 +39,7 @@ opentracing::expected<std::unique_ptr<opentracing::SpanContext>> ExtractImpl(
         make_error_code(std::errc::not_enough_memory));
   }
   std::unique_ptr<opentracing::SpanContext> span_context(mock_span_context);
-  auto result = mock_span_context->Extract(reader);
+  auto result = mock_span_context->Extract(propagation_options, reader);
   if (!result) {
     return opentracing::make_unexpected(result.error());
   }
@@ -41,6 +49,10 @@ opentracing::expected<std::unique_ptr<opentracing::SpanContext>> ExtractImpl(
   return std::move(span_context);
 }
 
+MockTracer::MockTracer(MockTracerOptions&& options)
+    : recorder_{std::move(options.recorder)},
+      propagation_options_{std::move(options.propagation_options)} {}
+
 std::unique_ptr<Span> MockTracer::StartSpanWithOptions(
     string_view operation_name, const StartSpanOptions& options) const
     noexcept try {
@@ -59,32 +71,32 @@ void MockTracer::Close() noexcept {
 
 expected<void> MockTracer::Inject(const SpanContext& sc,
                                   std::ostream& writer) const {
-  return InjectImpl(sc, writer);
+  return InjectImpl(propagation_options_, sc, writer);
 }
 
 expected<void> MockTracer::Inject(const SpanContext& sc,
                                   const TextMapWriter& writer) const {
-  return InjectImpl(sc, writer);
+  return InjectImpl(propagation_options_, sc, writer);
 }
 
 expected<void> MockTracer::Inject(const SpanContext& sc,
                                   const HTTPHeadersWriter& writer) const {
-  return InjectImpl(sc, writer);
+  return InjectImpl(propagation_options_, sc, writer);
 }
 
 expected<std::unique_ptr<SpanContext>> MockTracer::Extract(
     std::istream& reader) const {
-  return ExtractImpl(reader);
+  return ExtractImpl(propagation_options_, reader);
 }
 
 expected<std::unique_ptr<SpanContext>> MockTracer::Extract(
     const TextMapReader& reader) const {
-  return ExtractImpl(reader);
+  return ExtractImpl(propagation_options_, reader);
 }
 
 expected<std::unique_ptr<SpanContext>> MockTracer::Extract(
     const HTTPHeadersReader& reader) const {
-  return ExtractImpl(reader);
+  return ExtractImpl(propagation_options_, reader);
 }
 }  // namespace mocktracer
 END_OPENTRACING_ABI_NAMESPACE
diff --git a/mocktracer/test/propagation_test.cpp b/mocktracer/test/propagation_test.cpp
@@ -70,8 +70,10 @@ struct HTTPHeadersCarrier : HTTPHeadersReader, HTTPHeadersWriter {
 };
 
 TEST_CASE("propagation") {
+  const char* propagation_key = "propagation-key";
   auto recorder = new InMemoryRecorder{};
   MockTracerOptions tracer_options;
+  tracer_options.propagation_options.propagation_key = propagation_key;
   tracer_options.recorder.reset(recorder);
   auto tracer = std::shared_ptr<opentracing::Tracer>{
       new MockTracer{std::move(tracer_options)}};
@@ -82,6 +84,11 @@ TEST_CASE("propagation") {
   CHECK(span);
   span->SetBaggageItem("abc", "123");
 
+  SECTION("Propagation uses the specified propagation_key.") {
+    CHECK(tracer->Inject(span->context(), text_map_carrier));
+    CHECK(text_map.count(propagation_key) == 1);
+  }
+
   SECTION("Inject, extract, inject yields the same text_map.") {
     CHECK(tracer->Inject(span->context(), text_map_carrier));
     auto injection_map1 = text_map;
@@ -200,4 +207,30 @@ TEST_CASE("propagation") {
     CHECK(std::all_of(value.begin(), value.end(), is_base64_char));
     CHECK(value.size() % 4 == 0);
   }
+
+  SECTION("Inject fails if inject_error_code is non-zero.") {
+    MockTracerOptions tracer_options_fail;
+    auto error_code = std::make_error_code(std::errc::network_down);
+    tracer_options_fail.propagation_options.inject_error_code = error_code;
+    tracer = std::shared_ptr<opentracing::Tracer>{
+        new MockTracer{std::move(tracer_options_fail)}};
+
+    std::ostringstream oss;
+    auto rcode = tracer->Inject(span->context(), oss);
+    CHECK(!rcode);
+    CHECK(rcode.error() == error_code);
+  }
+
+  SECTION("Extract fails if extract_error_code is non-zero.") {
+    MockTracerOptions tracer_options_fail;
+    auto error_code = std::make_error_code(std::errc::network_down);
+    tracer_options_fail.propagation_options.extract_error_code = error_code;
+    tracer = std::shared_ptr<opentracing::Tracer>{
+        new MockTracer{std::move(tracer_options_fail)}};
+
+    CHECK(tracer->Inject(span->context(), text_map_carrier));
+    auto span_context_maybe = tracer->Extract(text_map_carrier);
+    CHECK(!span_context_maybe);
+    CHECK(span_context_maybe.error() == error_code);
+  }
 }
