From 98bd4f88db15f1a8598459db82b5c56b4a9d0005 Mon Sep 17 00:00:00 2001
From: saadiaelf <saadiaelfekak@gmail.com>
Date: Fri, 24 Jan 2025 07:38:44 -0800
Subject: [PATCH] add maxRequest test

---
 .../clusters/aggregate/cluster_test.cc        | 80 +++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/test/extensions/clusters/aggregate/cluster_test.cc b/test/extensions/clusters/aggregate/cluster_test.cc
index ec829deba4514..271489c70ef52 100644
--- a/test/extensions/clusters/aggregate/cluster_test.cc
+++ b/test/extensions/clusters/aggregate/cluster_test.cc
@@ -123,6 +123,9 @@ class AggregateClusterTest : public Event::TestUsingSimulatedTime, public testin
 
     setupPrioritySet();
 
+    EXPECT_CALL(*primary_info_, resourceManager(Upstream::ResourcePriority::Default)).Times(0);
+    EXPECT_CALL(*secondary_info_, resourceManager(Upstream::ResourcePriority::Default)).Times(0);
+
     ON_CALL(primary_, loadBalancer()).WillByDefault(ReturnRef(primary_load_balancer_));
     ON_CALL(secondary_, loadBalancer()).WillByDefault(ReturnRef(secondary_load_balancer_));
 
@@ -131,6 +134,16 @@ class AggregateClusterTest : public Event::TestUsingSimulatedTime, public testin
     lb_ = lb_factory_->create(lb_params_);
   }
 
+  // NOTE: in order to use remaining_<name> (and for it to not always be 0)
+  // we need to use "track_remaining: true" in the config to get the remaining_<name> stat
+  Stats::Gauge& getCircuitBreakersStatByPriority(std::string priority, std::string stat) {
+    std::string stat_name_ = "circuit_breakers." + priority + "." + stat;
+    Stats::StatNameManagedStorage statStore(stat_name_,
+                                            cluster_->info()->statsScope().symbolTable());
+    return cluster_->info()->statsScope().gaugeFromStatName(statStore.statName(),
+                                                            Stats::Gauge::ImportMode::Accumulate);
+  }
+
   NiceMock<Server::Configuration::MockServerFactoryContext> server_context_;
   Ssl::MockContextManager ssl_context_manager_;
 
@@ -169,6 +182,73 @@ class AggregateClusterTest : public Event::TestUsingSimulatedTime, public testin
 )EOF";
 }; // namespace Aggregate
 
+TEST_F(AggregateClusterTest, CircuitBreakerMaxRequestsTest) {
+  const std::string yaml_config = R"EOF(
+    name: aggregate_cluster
+    connect_timeout: 0.25s
+    lb_policy: CLUSTER_PROVIDED
+    circuit_breakers:
+      thresholds:
+      - priority: DEFAULT
+        max_requests: 1
+        track_remaining: true
+    cluster_type:
+      name: envoy.clusters.aggregate
+      typed_config:
+        "@type": type.googleapis.com/envoy.extensions.clusters.aggregate.v3.ClusterConfig
+        clusters:
+        - primary
+        - secondary
+)EOF";
+
+  initialize(yaml_config);
+
+  // resource manager for the DEFAULT priority (look above^)
+  Upstream::ResourceManager& resource_manager =
+      cluster_->info()->resourceManager(Upstream::ResourcePriority::Default);
+
+  // get the circuit breaker stats we are interested in, to assert against
+  Stats::Gauge& rq_open = getCircuitBreakersStatByPriority("default", "rq_open");
+  Stats::Gauge& remaining_rq = getCircuitBreakersStatByPriority("default", "remaining_rq");
+
+  // check the yaml config is set correctly
+  // we should have a maximum of 1 request available to use
+  EXPECT_EQ(1U, resource_manager.requests().max());
+
+  // check that we can create a new request
+  EXPECT_TRUE(resource_manager.requests().canCreate());
+  // check the requests count is 0
+  EXPECT_EQ(0U, resource_manager.requests().count());
+  // check that we have 1 remaining request
+  EXPECT_EQ(1U, remaining_rq.value());
+  // check the circuit breaker is closed
+  EXPECT_EQ(0U, rq_open.value());
+
+  // create that one request
+  resource_manager.requests().inc();
+
+  // check the request count is now 1
+  EXPECT_EQ(1U, resource_manager.requests().count());
+  // make sure we are NOT allowed to create anymore request
+  EXPECT_FALSE(resource_manager.requests().canCreate());
+  // check that we have 0 remaining requests
+  EXPECT_EQ(0U, remaining_rq.value());
+  // check the circuit breaker is now open
+  EXPECT_EQ(1U, rq_open.value());
+
+  // remove that one request
+  resource_manager.requests().dec();
+
+  // check the request count is now 0 again
+  EXPECT_EQ(0U, resource_manager.requests().count());
+  // check that we can create a new request again
+  EXPECT_TRUE(resource_manager.requests().canCreate());
+  // check that we have 1 remaining request again
+  EXPECT_EQ(1U, remaining_rq.value());
+  // check that the circuit breaker is closed again
+  EXPECT_EQ(0U, rq_open.value());
+}
+
 TEST_F(AggregateClusterTest, LoadBalancerTest) {
   initialize(default_yaml_config_);
   // Health value: