Merge pull request #105 from alexb5dh/main

Allow service overrides
Pier-Two · Jun 25, 2024 · 2f52bb0 · 2f52bb0
2 parents 59e8620 + 2639a21
commit 2f52bb0
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 14 deletions.
diff --git a/src/Lantern.Discv5.WireProtocol/Discv5ProtocolBuilder.cs b/src/Lantern.Discv5.WireProtocol/Discv5ProtocolBuilder.cs
@@ -18,6 +18,7 @@ public class Discv5ProtocolBuilder(IServiceCollection services) : IDiscv5Protoco
     private IEnrEntryRegistry _entryRegistry = new EnrEntryRegistry();
     private ILoggerFactory _loggerFactory = LoggingOptions.Default;
     private ITalkReqAndRespHandler? _talkResponder;
+    private Action<IServiceCollection>? _servicesSetup;
     private IServiceProvider? _serviceProvider;
 
     public IDiscv5ProtocolBuilder WithConnectionOptions(ConnectionOptions connectionOptions)
@@ -98,6 +99,12 @@ public IDiscv5ProtocolBuilder WithLoggerFactory(ILoggerFactory loggerFactory)
         return this;
     }
 
+    public IDiscv5ProtocolBuilder WithServices(Action<IServiceCollection> setup)
+    {
+        _servicesSetup = setup ?? throw new ArgumentNullException(nameof(setup));
+        return this;
+    }
+
     public IServiceProvider GetServiceProvider()
     {
         return _serviceProvider ?? throw new InvalidOperationException("Build() must be called before accessing the service provider.");
@@ -106,6 +113,7 @@ public IServiceProvider GetServiceProvider()
     public IDiscv5Protocol Build()
     {
         services.AddDiscv5(_tableOptions, _connectionOptions, _sessionOptions, _entryRegistry, _enrBuilder.Build(), _loggerFactory, _talkResponder);
+        _servicesSetup?.Invoke(services);
         _serviceProvider = services.BuildServiceProvider();
         return _serviceProvider.GetRequiredService<IDiscv5Protocol>();
     }

diff --git a/src/Lantern.Discv5.WireProtocol/IDiscv5ProtocolBuilder.cs b/src/Lantern.Discv5.WireProtocol/IDiscv5ProtocolBuilder.cs
@@ -36,6 +36,8 @@ public interface IDiscv5ProtocolBuilder
 
     IDiscv5ProtocolBuilder WithTalkResponder(Action<ITalkReqAndRespHandler> configure);
 
+    IDiscv5ProtocolBuilder WithServices(Action<IServiceCollection> setup);
+
     IServiceProvider GetServiceProvider();
 
     IDiscv5Protocol Build();

diff --git a/tests/Lantern.Discv5.WireProtocol.Tests/Discv5ProtocolBuilderTests.cs b/tests/Lantern.Discv5.WireProtocol.Tests/Discv5ProtocolBuilderTests.cs
@@ -134,25 +134,28 @@ public void WithTalkResponder_ChainsCorrectly_ReturnsBuilderInstance()
     [Test]
     public void Build_WithConfigurations_ReturnsConfiguredInstance()
     {
-        string[] bootstrapEnrs = ["enr:-example"];
         var services = new ServiceCollection();
         var discv5Builder = new Discv5ProtocolBuilder(services);
-        var sessionOptions = SessionOptions.Default;
-        var enr = new EnrBuilder()
-            .WithIdentityScheme(sessionOptions.Verifier, sessionOptions.Signer)
-            .WithEntry(EnrEntryKey.Id, new EntryId("v4"))
-            .WithEntry(EnrEntryKey.Secp256K1, new EntrySecp256K1(sessionOptions.Signer.PublicKey));
+        var discv5Protocol = AddMandatoryConfigurations(discv5Builder).Build();
 
-        var discv5Protocol = discv5Builder.WithConnectionOptions(new ConnectionOptions
-        {
-            UdpPort = 30303
-        }).WithTableOptions(new TableOptions(bootstrapEnrs))
-            .WithLoggerFactory(NullLoggerFactory.Instance)
-            .WithEnrBuilder(enr)
+        Assert.IsNotNull(discv5Protocol, "Expected to return a configured instance.");
+    }
+
+    [Test]
+    public void Build_WithServiceOverrides_UsesOverrides()
+    {
+        var services = new ServiceCollection();
+        var discv5Builder = new Discv5ProtocolBuilder(services);
+
+        var serviceOverride = Mock.Of<IConnectionManager>();
+        AddMandatoryConfigurations(discv5Builder, 30304)
+            .WithServices(s => s.AddSingleton(serviceOverride))
             .Build();
 
+        var serviceProvider = discv5Builder.GetServiceProvider();
+        var serviceResolved = serviceProvider.GetRequiredService<IConnectionManager>();
 
-        Assert.IsNotNull(discv5Protocol, "Expected to return a configured instance.");
+        Assert.AreSame(serviceOverride, serviceResolved, "Expected to return overriden service.");
     }
 
     [Test]
@@ -161,4 +164,20 @@ public void Build_WithoutMandatoryConfigurations_ThrowsException()
         _builder = new Discv5ProtocolBuilder(new ServiceCollection());
         Assert.Throws<InvalidOperationException>(() => _builder.Build(), "Expected to throw due to missing configurations.");
     }
-}
+
+    private static IDiscv5ProtocolBuilder AddMandatoryConfigurations(IDiscv5ProtocolBuilder discv5Builder, int udpPort = 30303)
+    {
+        string[] bootstrapEnrs = ["enr:-example"];
+        var sessionOptions = SessionOptions.Default;
+        var enr = new EnrBuilder()
+            .WithIdentityScheme(sessionOptions.Verifier, sessionOptions.Signer)
+            .WithEntry(EnrEntryKey.Id, new EntryId("v4"))
+            .WithEntry(EnrEntryKey.Secp256K1, new EntrySecp256K1(sessionOptions.Signer.PublicKey));
+
+        return discv5Builder
+            .WithConnectionOptions(new ConnectionOptions { UdpPort = udpPort })
+            .WithTableOptions(new TableOptions(bootstrapEnrs))
+            .WithLoggerFactory(NullLoggerFactory.Instance)
+            .WithEnrBuilder(enr);
+    }
+}