From fd35ad0514bc9875cde91245066b4e2400767553 Mon Sep 17 00:00:00 2001
From: David Lievrouw <david.lievrouw@recomatics.com>
Date: Wed, 10 Jul 2019 20:59:04 +0200
Subject: [PATCH] Add ability to register custom load balancers

---
 .../DependencyInjection/OcelotBuilderTests.cs | 44 ++++++++++++++++---
 1 file changed, 37 insertions(+), 7 deletions(-)

diff --git a/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs b/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
index f85bbb11..577fd059 100644
--- a/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
+++ b/test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs
@@ -154,6 +154,42 @@ namespace Ocelot.UnitTests.DependencyInjection
                 .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 @@ namespace Ocelot.UnitTests.DependencyInjection
         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 @@ namespace Ocelot.UnitTests.DependencyInjection
             _ocelotBuilder.AddSingletonDefinedAggregator<T>();
         }
         
-        private void AddCustomLoadBalancer<T>()
-            where T : ILoadBalancer, new()
-        {
-            _ocelotBuilder.AddCustomLoadBalancer<T>();
-        }
-
         private void AddTransientDefinedAggregator<T>()
             where T : class, IDefinedAggregator
         {