From 7490709db42e28d3e19212d4f27c50b3fe46208a Mon Sep 17 00:00:00 2001
From: Victor Vazquez <victor.vazquez@microsoft.com>
Date: Tue, 30 Mar 2021 11:52:47 -0700
Subject: [PATCH] test for Context heap integrity (#1999)

---
 sdk/core/azure-core/test/ut/context.cpp | 53 +++++++++++++++++++++++++
 sdk/core/azure-core/test/ut/policy.cpp  | 48 ++++++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/sdk/core/azure-core/test/ut/context.cpp b/sdk/core/azure-core/test/ut/context.cpp
index 2807d4fc2..debb7a318 100644
--- a/sdk/core/azure-core/test/ut/context.cpp
+++ b/sdk/core/azure-core/test/ut/context.cpp
@@ -191,3 +191,56 @@ TEST(Context, UniquePtr)
   auto& contextValueRef = contextP.Get<std::unique_ptr<SomeStructForContext>>("struct");
   EXPECT_EQ(contextValueRef->someField, 12345);
 }
+
+TEST(Context, HeapLinkIntegrity)
+{
+  Context thirdGeneration; // To be used at the end
+  {
+    Context root;
+    auto firstGeneration = root.WithValue("a", std::string("a"));
+    EXPECT_TRUE(firstGeneration.HasKey("a"));
+
+    auto secondGeneration = firstGeneration.WithValue("b", std::string("b"));
+    EXPECT_TRUE(secondGeneration.HasKey("a"));
+    EXPECT_EQ("a", secondGeneration.Get<std::string>("a"));
+    EXPECT_TRUE(secondGeneration.HasKey("b"));
+    EXPECT_EQ("b", secondGeneration.Get<std::string>("b"));
+
+    // Now overide the generation
+    secondGeneration = secondGeneration.WithValue("c", std::string("c"));
+    EXPECT_TRUE(
+        secondGeneration.HasKey("a")); // Still know about first gen - The link is still in heap
+    EXPECT_EQ("a", secondGeneration.Get<std::string>("a"));
+    EXPECT_TRUE(secondGeneration.HasKey(
+        "b")); // Still knows about the initial second gen, as a shared_ptr, it is still on heap
+    EXPECT_EQ("b", secondGeneration.Get<std::string>("b"));
+    EXPECT_TRUE(secondGeneration.HasKey("c")); // Check new value
+    EXPECT_EQ("c", secondGeneration.Get<std::string>("c"));
+
+    // One more override
+    secondGeneration = secondGeneration.WithValue("d", std::string("d"));
+    EXPECT_TRUE(secondGeneration.HasKey("a"));
+    EXPECT_EQ("a", secondGeneration.Get<std::string>("a"));
+    EXPECT_TRUE(secondGeneration.HasKey("b"));
+    EXPECT_EQ("b", secondGeneration.Get<std::string>("b"));
+    EXPECT_TRUE(secondGeneration.HasKey("c"));
+    EXPECT_EQ("c", secondGeneration.Get<std::string>("c"));
+    EXPECT_TRUE(secondGeneration.HasKey("d"));
+    EXPECT_EQ("d", secondGeneration.Get<std::string>("d"));
+
+    // New Gen
+    thirdGeneration = secondGeneration.WithValue("e", std::string("e"));
+  }
+  // Went out of scope, root and secondGeneration are destroyed. but should remain in heap for the
+  // third-generation since the previous geneations are still alive inside his heart <3.
+  EXPECT_TRUE(thirdGeneration.HasKey("a"));
+  EXPECT_EQ("a", thirdGeneration.Get<std::string>("a"));
+  EXPECT_TRUE(thirdGeneration.HasKey("b"));
+  EXPECT_EQ("b", thirdGeneration.Get<std::string>("b"));
+  EXPECT_TRUE(thirdGeneration.HasKey("c"));
+  EXPECT_EQ("c", thirdGeneration.Get<std::string>("c"));
+  EXPECT_TRUE(thirdGeneration.HasKey("d"));
+  EXPECT_EQ("d", thirdGeneration.Get<std::string>("d"));
+  EXPECT_TRUE(thirdGeneration.HasKey("e"));
+  EXPECT_EQ("e", thirdGeneration.Get<std::string>("e"));
+}
diff --git a/sdk/core/azure-core/test/ut/policy.cpp b/sdk/core/azure-core/test/ut/policy.cpp
index 148430bcb..1c36fc579 100644
--- a/sdk/core/azure-core/test/ut/policy.cpp
+++ b/sdk/core/azure-core/test/ut/policy.cpp
@@ -44,6 +44,28 @@ struct TestRetryPolicySharedState : public Azure::Core::Http::Policies::HttpPoli
   }
 };
 
+struct TestContextTreeIntegrity : public Azure::Core::Http::Policies::HttpPolicy
+{
+  std::unique_ptr<HttpPolicy> Clone() const override
+  {
+    return std::make_unique<TestContextTreeIntegrity>(*this);
+  }
+
+  std::unique_ptr<Azure::Core::Http::RawResponse> Send(
+      Azure::Core::Http::Request& request,
+      Azure::Core::Http::Policies::NextHttpPolicy nextHttpPolicy,
+      Azure::Core::Context const& ctx) const override
+  {
+    EXPECT_TRUE(ctx.HasKey("TheKey"));
+    if (ctx.HasKey("TheKey"))
+    {
+      auto value = ctx.Get<std::string>("TheKey");
+      EXPECT_EQ("TheValue", value);
+    }
+    return nextHttpPolicy.Send(request, ctx);
+  }
+};
+
 class SuccessAfter : public Azure::Core::Http::Policies::HttpPolicy {
 private:
   int m_successAfter; // Always success
@@ -172,3 +194,29 @@ TEST(Policy, RetryPolicyRetryCycle)
   Request request(HttpMethod::Get, Url("url"));
   pipeline.Send(request, Context::GetApplicationContext());
 }
+
+// Makes sure that the context tree is not corrupted/broken by some policy
+TEST(Policy, RetryPolicyKeepContext)
+{
+  using namespace Azure::Core;
+  using namespace Azure::Core::Http;
+  using namespace Azure::Core::Http::Policies;
+  using namespace Azure::Core::Http::_internal;
+  // Clean the validation global state
+  retryCounterState = 0;
+
+  // Pipeline with retry test
+  std::vector<std::unique_ptr<HttpPolicy>> policies;
+  RetryOptions opt;
+  opt.RetryDelay = std::chrono::milliseconds(10);
+  policies.push_back(std::make_unique<RetryPolicy>(opt));
+  policies.push_back(std::make_unique<TestRetryPolicySharedState>());
+  policies.push_back(std::make_unique<TestContextTreeIntegrity>());
+  policies.push_back(std::make_unique<SuccessAfter>(3));
+
+  HttpPipeline pipeline(policies);
+  Request request(HttpMethod::Get, Url("url"));
+  auto withValueContext
+      = Context::GetApplicationContext().WithValue("TheKey", std::string("TheValue"));
+  pipeline.Send(request, withValueContext);
+}