Add ability to register custom load balancers

ThreeMammals · Mar 15, 2020 · fd35ad0 · fd35ad0
1 parent 80ab5f6
commit fd35ad0
Showing 1 changed file with 37 additions and 7 deletions.
diff --git a/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs b/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
@@ -154,6 +154,42 @@ public void should_add_transient_defined_aggregators()
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_default_ctor()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer<FakeCustomLoadBalancer>())
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_factory_method()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer(() => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_di_factory_method()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer(provider => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_add_custom_load_balancer_creators_by_factory_method_with_arguments()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer((reroute, discoveryProvider) => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
         [Fact]
         public void should_replace_iplaceholder()
         {
@@ -168,7 +204,7 @@ public void should_replace_iplaceholder()
         public void should_add_custom_load_balancer_creators()
         {
             this.Given(x => WhenISetUpOcelotServices())
-                .When(x => AddCustomLoadBalancer<FakeCustomLoadBalancer>())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer((provider, reroute, discoveryProvider) => new FakeCustomLoadBalancer()))
                 .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
                 .BDDfy();
         }
@@ -179,12 +215,6 @@ private void AddSingletonDefinedAggregator<T>()
             _ocelotBuilder.AddSingletonDefinedAggregator<T>();
         }
 
-        private void AddCustomLoadBalancer<T>()
-            where T : ILoadBalancer, new()
-        {
-            _ocelotBuilder.AddCustomLoadBalancer<T>();
-        }
-
         private void AddTransientDefinedAggregator<T>()
             where T : class, IDefinedAggregator
         {