Merge branch 'CustomLoadBalancers' of https://github.com/DavidLievrouw/Ocelot into DavidLievrouw-CustomLoadBalancers

src/Ocelot/DependencyInjection/IOcelotBuilder.cs

@@ -3,6 +3,9 @@ using Microsoft.Extensions.DependencyInjection;
 using Ocelot.Middleware.Multiplexer;
 using System;
 using System.Net.Http;
+using Ocelot.Configuration;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.ServiceDiscovery.Providers;
 
 namespace Ocelot.DependencyInjection
 {
@@ -25,6 +28,23 @@ namespace Ocelot.DependencyInjection
         IOcelotBuilder AddTransientDefinedAggregator<T>()
             where T : class, IDefinedAggregator;
 
+        IOcelotBuilder AddCustomLoadBalancer<T>()
+            where T : ILoadBalancer, new();
+        
+        IOcelotBuilder AddCustomLoadBalancer<T>(Func<T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+        IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+        IOcelotBuilder AddCustomLoadBalancer<T>(
+            Func<DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
+        IOcelotBuilder AddCustomLoadBalancer<T>(
+            Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer;
+
         IOcelotBuilder AddConfigPlaceholders();
     }
 }

src/Ocelot/DependencyInjection/OcelotBuilder.cs

@@ -1,3 +1,5 @@
+using Ocelot.ServiceDiscovery.Providers;
+
 using Ocelot.Configuration.ChangeTracking;
 
 namespace Ocelot.DependencyInjection
@@ -87,6 +89,10 @@ namespace Ocelot.DependencyInjection
             Services.TryAddSingleton<IFileConfigurationRepository, DiskFileConfigurationRepository>();
             Services.TryAddSingleton<IFileConfigurationSetter, FileAndInternalConfigurationSetter>();
             Services.TryAddSingleton<IServiceDiscoveryProviderFactory, ServiceDiscoveryProviderFactory>();
+            Services.AddSingleton<ILoadBalancerCreator, NoLoadBalancerCreator>();
+            Services.AddSingleton<ILoadBalancerCreator, RoundRobinCreator>();
+            Services.AddSingleton<ILoadBalancerCreator, CookieStickySessionsCreator>();
+            Services.AddSingleton<ILoadBalancerCreator, LeastConnectionCreator>();
             Services.TryAddSingleton<ILoadBalancerFactory, LoadBalancerFactory>();
             Services.TryAddSingleton<ILoadBalancerHouse, LoadBalancerHouse>();
             Services.TryAddSingleton<IOcelotLoggerFactory, AspDotNetLoggerFactory>();
@@ -169,6 +175,47 @@ namespace Ocelot.DependencyInjection
             return this;
         }
 
+        public IOcelotBuilder AddCustomLoadBalancer<T>()
+            where T : ILoadBalancer, new()
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) => new T());
+            return this;
+        }
+        
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<T> loadBalancerFactoryFunc) 
+            where T : ILoadBalancer
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) =>
+                loadBalancerFactoryFunc());
+            return this;
+        }
+
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, T> loadBalancerFactoryFunc) 
+            where T : ILoadBalancer
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) =>
+                loadBalancerFactoryFunc(provider));
+            return this;
+        }
+
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer
+        {
+            AddCustomLoadBalancer((provider, reRoute, serviceDiscoveryProvider) =>
+                loadBalancerFactoryFunc(reRoute, serviceDiscoveryProvider));
+            return this;
+        }
+
+        public IOcelotBuilder AddCustomLoadBalancer<T>(Func<IServiceProvider, DownstreamReRoute, IServiceDiscoveryProvider, T> loadBalancerFactoryFunc)
+            where T : ILoadBalancer
+        {
+            Services.AddSingleton<ILoadBalancerCreator>(provider =>
+                new DelegateInvokingLoadBalancerCreator<T>(
+                    (reRoute, serviceDiscoveryProvider) => 
+                        loadBalancerFactoryFunc(provider, reRoute, serviceDiscoveryProvider)));
+            return this;
+        }
+
         private void AddSecurity()
         {
             Services.TryAddSingleton<ISecurityOptionsCreator, SecurityOptionsCreator>();

src/Ocelot/LoadBalancer/LoadBalancers/CookieStickySessionsCreator.cs

@@ -0,0 +1,20 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System.Threading.Tasks;
+    using Ocelot.Configuration;
+    using Ocelot.Infrastructure;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class CookieStickySessionsCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            var loadBalancer = new RoundRobin(async () => await serviceProvider.Get());
+            var bus = new InMemoryBus<StickySession>();
+            return new CookieStickySessions(loadBalancer, reRoute.LoadBalancerOptions.Key,
+                reRoute.LoadBalancerOptions.ExpiryInMs, bus);
+        }
+
+        public string Type => nameof(CookieStickySessions);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/DelegateInvokingLoadBalancerCreator.cs

@@ -0,0 +1,25 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using System;
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class DelegateInvokingLoadBalancerCreator<T> : ILoadBalancerCreator
+        where T : ILoadBalancer
+    {
+        private readonly Func<DownstreamReRoute, IServiceDiscoveryProvider, ILoadBalancer> _creatorFunc;
+
+        public DelegateInvokingLoadBalancerCreator(
+            Func<DownstreamReRoute, IServiceDiscoveryProvider, ILoadBalancer> creatorFunc)
+        {
+            _creatorFunc = creatorFunc;
+        }
+
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return _creatorFunc(reRoute, serviceProvider);
+        }
+
+        public string Type => typeof(T).Name;
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/ILoadBalancerCreator.cs

@@ -0,0 +1,11 @@
+using Ocelot.Configuration;
+using Ocelot.ServiceDiscovery.Providers;
+
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    public interface ILoadBalancerCreator
+    {
+        ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider);
+        string Type { get; }
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/LeastConnectionCreator.cs

@@ -0,0 +1,15 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class LeastConnectionCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new LeastConnection(async () => await serviceProvider.Get(), reRoute.ServiceName);
+        }
+
+        public string Type => nameof(LeastConnection);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/LoadBalancerFactory.cs

@@ -1,47 +1,42 @@
-using Ocelot.Configuration;
+namespace Ocelot.LoadBalancer.LoadBalancers
-using Ocelot.Infrastructure;
-using Ocelot.Responses;
-using Ocelot.ServiceDiscovery;
-using System.Threading.Tasks;
-
-namespace Ocelot.LoadBalancer.LoadBalancers
 {
+    using System.Collections.Generic;
+    using System.Linq;
+    using Ocelot.Configuration;
+    using Ocelot.Responses;
+    using System.Threading.Tasks;
+    using Ocelot.ServiceDiscovery;
+
     public class LoadBalancerFactory : ILoadBalancerFactory
     {
         private readonly IServiceDiscoveryProviderFactory _serviceProviderFactory;
+        private readonly IEnumerable<ILoadBalancerCreator> _loadBalancerCreators;
 
-        public LoadBalancerFactory(IServiceDiscoveryProviderFactory serviceProviderFactory)
+        public LoadBalancerFactory(IServiceDiscoveryProviderFactory serviceProviderFactory, IEnumerable<ILoadBalancerCreator> loadBalancerCreators)
         {
             _serviceProviderFactory = serviceProviderFactory;
+            _loadBalancerCreators = loadBalancerCreators;
         }
 
-        public async Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
+        public Task<Response<ILoadBalancer>> Get(DownstreamReRoute reRoute, ServiceProviderConfiguration config)
         {
-            var response = _serviceProviderFactory.Get(config, reRoute);
+            var serviceProviderFactoryResponse = _serviceProviderFactory.Get(config, reRoute);
 
-            if (response.IsError)
+            Response<ILoadBalancer> response;
+            if (serviceProviderFactoryResponse.IsError)
             {
-                return new ErrorResponse<ILoadBalancer>(response.Errors);
+                response = new ErrorResponse<ILoadBalancer>(serviceProviderFactoryResponse.Errors);
             }
-
+            else
-            var serviceProvider = response.Data;
-
-            switch (reRoute.LoadBalancerOptions?.Type)
             {
-                case nameof(RoundRobin):
+                var serviceProvider = serviceProviderFactoryResponse.Data;
-                    return new OkResponse<ILoadBalancer>(new RoundRobin(async () => await serviceProvider.Get()));
+                var requestedType = reRoute.LoadBalancerOptions?.Type ?? nameof(NoLoadBalancer);
+                var applicableCreator = _loadBalancerCreators.Single(c => c.Type == requestedType);
+                var createdLoadBalancer = applicableCreator.Create(reRoute, serviceProvider);
+                response = new OkResponse<ILoadBalancer>(createdLoadBalancer);
+            }
 
-                case nameof(LeastConnection):
+            return Task.FromResult(response);
-                    return new OkResponse<ILoadBalancer>(new LeastConnection(async () => await serviceProvider.Get(), reRoute.ServiceName));
-
-                case nameof(CookieStickySessions):
-                    var loadBalancer = new RoundRobin(async () => await serviceProvider.Get());
-                    var bus = new InMemoryBus<StickySession>();
-                    return new OkResponse<ILoadBalancer>(new CookieStickySessions(loadBalancer, reRoute.LoadBalancerOptions.Key, reRoute.LoadBalancerOptions.ExpiryInMs, bus));
-
-                default:
-                    return new OkResponse<ILoadBalancer>(new NoLoadBalancer(async () => await serviceProvider.Get()));
-            }
         }
     }
 }

src/Ocelot/LoadBalancer/LoadBalancers/NoLoadBalancerCreator.cs

@@ -0,0 +1,15 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{ 
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+
+    public class NoLoadBalancerCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new NoLoadBalancer(async () => await serviceProvider.Get());
+        }
+
+        public string Type => nameof(NoLoadBalancer);
+    }
+}

src/Ocelot/LoadBalancer/LoadBalancers/RoundRobin.cs

@@ -1,12 +1,12 @@
-using Ocelot.Middleware;
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Ocelot.Middleware;
     using Ocelot.Responses;
     using Ocelot.Values;
     using System;
     using System.Collections.Generic;
     using System.Threading.Tasks;
 
-namespace Ocelot.LoadBalancer.LoadBalancers
-{
     public class RoundRobin : ILoadBalancer
     {
         private readonly Func<Task<List<Service>>> _services;

src/Ocelot/LoadBalancer/LoadBalancers/RoundRobinCreator.cs

@@ -0,0 +1,15 @@
+namespace Ocelot.LoadBalancer.LoadBalancers
+{
+    using Ocelot.Configuration;
+    using Ocelot.ServiceDiscovery.Providers;
+    
+    public class RoundRobinCreator : ILoadBalancerCreator
+    {
+        public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+        {
+            return new RoundRobin(async () => await serviceProvider.Get());
+        }
+
+        public string Type => nameof(RoundRobin);
+    }
+}

test/Ocelot.UnitTests/DependencyInjection/OcelotBuilderTests.cs

@@ -1,3 +1,9 @@
+using System.Threading.Tasks;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.Middleware;
+using Ocelot.Responses;
+using Ocelot.Values;
+
 namespace Ocelot.UnitTests.DependencyInjection
 {
     using Microsoft.AspNetCore.Hosting;
@@ -148,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()
         {
@@ -158,6 +200,15 @@ namespace Ocelot.UnitTests.DependencyInjection
                 .BDDfy();
         }
 
+        [Fact]
+        public void should_add_custom_load_balancer_creators()
+        {
+            this.Given(x => WhenISetUpOcelotServices())
+                .When(x => _ocelotBuilder.AddCustomLoadBalancer((provider, reroute, discoveryProvider) => new FakeCustomLoadBalancer()))
+                .Then(x => ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators())
+                .BDDfy();
+        }
+
         private void AddSingletonDefinedAggregator<T>()
             where T : class, IDefinedAggregator
         {
@@ -239,6 +290,17 @@ namespace Ocelot.UnitTests.DependencyInjection
             handlers[1].ShouldBeOfType<TWo>();
         }
         
+        private void ThenTheProviderIsRegisteredAndReturnsBothBuiltInAndCustomLoadBalancerCreators()
+        {
+            _serviceProvider = _services.BuildServiceProvider();
+            var creators = _serviceProvider.GetServices<ILoadBalancerCreator>().ToList();
+            creators.Count(c => c.GetType() == typeof(NoLoadBalancerCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(RoundRobinCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(CookieStickySessionsCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(LeastConnectionCreator)).ShouldBe(1);
+            creators.Count(c => c.GetType() == typeof(DelegateInvokingLoadBalancerCreator<FakeCustomLoadBalancer>)).ShouldBe(1);
+        }
+
         private void ThenTheAggregatorsAreTransient<TOne, TWo>()
         {
             var aggregators = _serviceProvider.GetServices<IDefinedAggregator>().ToList();
@@ -323,5 +385,20 @@ namespace Ocelot.UnitTests.DependencyInjection
         {
             _ex.ShouldBeNull();
         }
+
+        private class FakeCustomLoadBalancer : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                // Not relevant for these tests
+                throw new NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                // Not relevant for these tests
+                throw new NotImplementedException();
+            }
+        }
     }
 }

test/Ocelot.UnitTests/LoadBalancer/CookieStickySessionsCreatorTests.cs

@@ -0,0 +1,73 @@
+using Moq;
+using Ocelot.Configuration;
+using Ocelot.Configuration.Builder;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.ServiceDiscovery.Providers;
+using Shouldly;
+using TestStack.BDDfy;
+using Xunit;
+
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    public class CookieStickySessionsCreatorTests
+    {
+        private readonly CookieStickySessionsCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private ILoadBalancer _loadBalancer;
+        private string _typeName;
+
+        public CookieStickySessionsCreatorTests()
+        {
+            _creator = new CookieStickySessionsCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .WithLoadBalancerOptions(new LoadBalancerOptions("myType", "myKey", 1000))
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<CookieStickySessions>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("CookieStickySessions"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/DelegateInvokingLoadBalancerCreatorTests.cs

@@ -0,0 +1,102 @@
+using System;
+using System.Threading.Tasks;
+using Moq;
+using Ocelot.Configuration;
+using Ocelot.Configuration.Builder;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.Middleware;
+using Ocelot.Responses;
+using Ocelot.ServiceDiscovery.Providers;
+using Ocelot.Values;
+using Shouldly;
+using TestStack.BDDfy;
+using Xunit;
+
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    public class DelegateInvokingLoadBalancerCreatorTests
+    {
+        private readonly DelegateInvokingLoadBalancerCreator<FakeLoadBalancer> _creator;
+        private readonly Func<DownstreamReRoute, IServiceDiscoveryProvider, ILoadBalancer> _creatorFunc;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private ILoadBalancer _loadBalancer;
+        private string _typeName;
+
+        public DelegateInvokingLoadBalancerCreatorTests()
+        {
+            _creatorFunc = (reRoute, serviceDiscoveryProvider) =>
+                new FakeLoadBalancer(reRoute, serviceDiscoveryProvider);
+            _creator = new DelegateInvokingLoadBalancerCreator<FakeLoadBalancer>(_creatorFunc);
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("FakeLoadBalancer"))
+                .BDDfy();
+        }
+
+        [Fact]
+        public void should_return_result_of_specified_creator_func()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<FakeLoadBalancer>())
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+
+        private class FakeLoadBalancer : ILoadBalancer
+        {
+            public FakeLoadBalancer(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceDiscoveryProvider)
+            {
+                ReRoute = reRoute;
+                ServiceDiscoveryProvider = serviceDiscoveryProvider;
+            }
+
+            public DownstreamReRoute ReRoute { get; }
+            public IServiceDiscoveryProvider ServiceDiscoveryProvider { get; }
+
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new NotImplementedException();
+            }
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/LeastConnectionCreatorTests.cs

@@ -0,0 +1,73 @@
+using Moq;
+using Ocelot.Configuration;
+using Ocelot.Configuration.Builder;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.ServiceDiscovery.Providers;
+using Shouldly;
+using TestStack.BDDfy;
+using Xunit;
+
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    public class LeastConnectionCreatorTests
+    {
+        private readonly LeastConnectionCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private ILoadBalancer _loadBalancer;
+        private string _typeName;
+
+        public LeastConnectionCreatorTests()
+        {
+            _creator = new LeastConnectionCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .WithServiceName("myService")
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<LeastConnection>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("LeastConnection"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/LoadBalancerFactoryTests.cs

@@ -7,6 +7,10 @@ using Ocelot.ServiceDiscovery;
 using Ocelot.ServiceDiscovery.Providers;
 using Shouldly;
 using System.Collections.Generic;
+using System.Threading.Tasks;
+using Ocelot.Infrastructure.RequestData;
+using Ocelot.Middleware;
+using Ocelot.Values;
 using TestStack.BDDfy;
 using Xunit;
 
@@ -18,6 +22,7 @@ namespace Ocelot.UnitTests.LoadBalancer
         private readonly LoadBalancerFactory _factory;
         private Response<ILoadBalancer> _result;
         private readonly Mock<IServiceDiscoveryProviderFactory> _serviceProviderFactory;
+        private readonly IEnumerable<ILoadBalancerCreator> _loadBalancerCreators;
         private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
         private ServiceProviderConfiguration _serviceProviderConfig;
 
@@ -25,11 +30,17 @@ namespace Ocelot.UnitTests.LoadBalancer
         {
             _serviceProviderFactory = new Mock<IServiceDiscoveryProviderFactory>();
             _serviceProvider = new Mock<IServiceDiscoveryProvider>();
-            _factory = new LoadBalancerFactory(_serviceProviderFactory.Object);
+            _loadBalancerCreators = new ILoadBalancerCreator[]
+            {
+                new FakeLoadBalancerCreator<FakeLoadBalancerOne>(),
+                new FakeLoadBalancerCreator<FakeLoadBalancerTwo>(),
+                new FakeLoadBalancerCreator<FakeNoLoadBalancer>(nameof(NoLoadBalancer)),
+            };
+            _factory = new LoadBalancerFactory(_serviceProviderFactory.Object, _loadBalancerCreators);
         }
 
         [Fact]
-        public void should_return_no_load_balancer()
+        public void should_return_no_load_balancer_by_default()
         {
             var reRoute = new DownstreamReRouteBuilder()
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
@@ -39,15 +50,15 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
                 .And(x => x.GivenTheServiceProviderFactoryReturns())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<NoLoadBalancer>())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<FakeNoLoadBalancer>())
                 .BDDfy();
         }
 
         [Fact]
-        public void should_return_round_robin_load_balancer()
+        public void should_return_matching_load_balancer()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("RoundRobin", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("FakeLoadBalancerTwo", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
@@ -55,23 +66,7 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
                 .And(x => x.GivenTheServiceProviderFactoryReturns())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<RoundRobin>())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<FakeLoadBalancerTwo>())
-                .BDDfy();
-        }
-
-        [Fact]
-        public void should_return_round_least_connection_balancer()
-        {
-            var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("LeastConnection", "", 0))
-                .WithUpstreamHttpMethod(new List<string> { "Get" })
-                .Build();
-
-            this.Given(x => x.GivenAReRoute(reRoute))
-                .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
-                .And(x => x.GivenTheServiceProviderFactoryReturns())
-                .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<LeastConnection>())
                 .BDDfy();
         }
         
@@ -79,7 +74,7 @@ namespace Ocelot.UnitTests.LoadBalancer
         public void should_call_service_provider()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("RoundRobin", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("FakeLoadBalancerOne", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
@@ -92,18 +87,18 @@ namespace Ocelot.UnitTests.LoadBalancer
         }
 
         [Fact]
-        public void should_return_sticky_session()
+        public void should_return_error_response_when_call_to_service_provider_fails()
         {
             var reRoute = new DownstreamReRouteBuilder()
-                .WithLoadBalancerOptions(new LoadBalancerOptions("CookieStickySessions", "", 0))
+                .WithLoadBalancerOptions(new LoadBalancerOptions("FakeLoadBalancerOne", "", 0))
                 .WithUpstreamHttpMethod(new List<string> { "Get" })
                 .Build();
 
             this.Given(x => x.GivenAReRoute(reRoute))
                 .And(x => GivenAServiceProviderConfig(new ServiceProviderConfigurationBuilder().Build()))
-                .And(x => x.GivenTheServiceProviderFactoryReturns())
+                .And(x => x.GivenTheServiceProviderFactoryFails())
                 .When(x => x.WhenIGetTheLoadBalancer())
-                .Then(x => x.ThenTheLoadBalancerIsReturned<CookieStickySessions>())
+                .Then(x => x.ThenAnErrorResponseIsReturned())
                 .BDDfy();
         }
 
@@ -119,6 +114,13 @@ namespace Ocelot.UnitTests.LoadBalancer
                 .Returns(new OkResponse<IServiceDiscoveryProvider>(_serviceProvider.Object));
         }
 
+        private void GivenTheServiceProviderFactoryFails()
+        {
+            _serviceProviderFactory
+                .Setup(x => x.Get(It.IsAny<ServiceProviderConfiguration>(), It.IsAny<DownstreamReRoute>()))
+                .Returns(new ErrorResponse<IServiceDiscoveryProvider>(new CannotFindDataError("For tests")));
+        }
+
         private void ThenTheServiceProviderIsCalledCorrectly()
         {
             _serviceProviderFactory
@@ -139,5 +141,71 @@ namespace Ocelot.UnitTests.LoadBalancer
         {
             _result.Data.ShouldBeOfType<T>();
         }
+
+        private void ThenAnErrorResponseIsReturned()
+        {
+            _result.IsError.ShouldBeTrue();
+        }
+
+        private class FakeLoadBalancerCreator<T> : ILoadBalancerCreator
+            where T : ILoadBalancer, new()
+        {
+
+            public FakeLoadBalancerCreator()
+            {
+                Type = typeof(T).Name;
+            }
+
+            public FakeLoadBalancerCreator(string type)
+            {
+                Type = type;
+            }
+
+            public ILoadBalancer Create(DownstreamReRoute reRoute, IServiceDiscoveryProvider serviceProvider)
+            {
+                return new T();
+            }
+            
+            public string Type { get; }
+        }
+
+        private class FakeLoadBalancerOne : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
+        private class FakeLoadBalancerTwo : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
+
+        private class FakeNoLoadBalancer : ILoadBalancer
+        {
+            public Task<Response<ServiceHostAndPort>> Lease(DownstreamContext context)
+            {
+                throw new System.NotImplementedException();
+            }
+
+            public void Release(ServiceHostAndPort hostAndPort)
+            {
+                throw new System.NotImplementedException();
+            }
+        }
     }
 }

test/Ocelot.UnitTests/LoadBalancer/NoLoadBalancerCreatorTests.cs

@@ -0,0 +1,72 @@
+using Moq;
+using Ocelot.Configuration;
+using Ocelot.Configuration.Builder;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.ServiceDiscovery.Providers;
+using Shouldly;
+using TestStack.BDDfy;
+using Xunit;
+
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    public class NoLoadBalancerCreatorTests
+    {
+        private readonly NoLoadBalancerCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private ILoadBalancer _loadBalancer;
+        private string _typeName;
+
+        public NoLoadBalancerCreatorTests()
+        {
+            _creator = new NoLoadBalancerCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<NoLoadBalancer>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("NoLoadBalancer"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}

test/Ocelot.UnitTests/LoadBalancer/RoundRobinCreatorTests.cs

@@ -0,0 +1,72 @@
+using Moq;
+using Ocelot.Configuration;
+using Ocelot.Configuration.Builder;
+using Ocelot.LoadBalancer.LoadBalancers;
+using Ocelot.ServiceDiscovery.Providers;
+using Shouldly;
+using TestStack.BDDfy;
+using Xunit;
+
+namespace Ocelot.UnitTests.LoadBalancer
+{
+    public class RoundRobinCreatorTests
+    {
+        private readonly RoundRobinCreator _creator;
+        private readonly Mock<IServiceDiscoveryProvider> _serviceProvider;
+        private DownstreamReRoute _reRoute;
+        private ILoadBalancer _loadBalancer;
+        private string _typeName;
+
+        public RoundRobinCreatorTests()
+        {
+            _creator = new RoundRobinCreator();
+            _serviceProvider = new Mock<IServiceDiscoveryProvider>();
+        }
+        
+        [Fact]
+        public void should_return_instance_of_expected_load_balancer_type()
+        {
+            var reRoute = new DownstreamReRouteBuilder()
+                .Build();
+
+            this.Given(x => x.GivenAReRoute(reRoute))
+                .When(x => x.WhenIGetTheLoadBalancer())
+                .Then(x => x.ThenTheLoadBalancerIsReturned<RoundRobin>())
+                .BDDfy();
+        }
+                
+        [Fact]
+        public void should_return_expected_name()
+        {
+            this.When(x => x.WhenIGetTheLoadBalancerTypeName())
+                .Then(x => x.ThenTheLoadBalancerTypeIs("RoundRobin"))
+                .BDDfy();
+        }
+
+        private void GivenAReRoute(DownstreamReRoute reRoute)
+        {
+            _reRoute = reRoute;
+        }
+
+        private void WhenIGetTheLoadBalancer()
+        {
+            _loadBalancer = _creator.Create(_reRoute, _serviceProvider.Object);
+        }
+        
+        private void WhenIGetTheLoadBalancerTypeName()
+        {
+            _typeName = _creator.Type;
+        }
+
+        private void ThenTheLoadBalancerIsReturned<T>()
+            where T : ILoadBalancer
+        {
+            _loadBalancer.ShouldBeOfType<T>();
+        }
+
+        private void ThenTheLoadBalancerTypeIs(string type)
+        {
+            _typeName.ShouldBe(type);
+        }
+    }
+}