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

2025-06-19 15:58:16 +08:00 · 2020-04-13 11:06:49 +01:00
parent d58f31475e 1223c1985f
commit ccb5b84103
17 changed files with 767 additions and 67 deletions
--- a/test/Ocelot.UnitTests/LoadBalancer/CookieStickySessionsCreatorTests.cs
+++ b/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);
+        }
+    }
+}
--- a/test/Ocelot.UnitTests/LoadBalancer/DelegateInvokingLoadBalancerCreatorTests.cs
+++ b/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();
+            }
+        }
+    }
+}
--- a/test/Ocelot.UnitTests/LoadBalancer/LeastConnectionCreatorTests.cs
+++ b/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);
+        }
+    }
+}
--- a/test/Ocelot.UnitTests/LoadBalancer/LoadBalancerFactoryTests.cs
+++ b/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,31 +66,15 @@ 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();
-        }
-
+        
        [Fact]
        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();
+            }
+        }
    }
 }
--- a/test/Ocelot.UnitTests/LoadBalancer/NoLoadBalancerCreatorTests.cs
+++ b/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);
+        }
+    }
+}
--- a/test/Ocelot.UnitTests/LoadBalancer/RoundRobinCreatorTests.cs
+++ b/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);
+        }
+    }
+}