Merge pull request #17 from Pier-Two/dev

Add NUnit tests
Pier-Two · Jul 30, 2024 · 826f7e7 · 826f7e7
2 parents 827a98f + 6bd98bd
commit 826f7e7
Show file tree

Hide file tree

Showing 12 changed files with 407 additions and 44 deletions.
diff --git a/src/Lantern.Beacon/BeaconClient.cs b/src/Lantern.Beacon/BeaconClient.cs
@@ -4,8 +4,6 @@
 using Lantern.Beacon.Networking.Gossip;
 using Lantern.Beacon.Storage;
 using Lantern.Beacon.Sync;
-using Lantern.Beacon.Sync.Helpers;
-using Lantern.Beacon.Sync.Processors;
 using Lantern.Beacon.Sync.Types.Ssz.Altair;
 using Lantern.Beacon.Sync.Types.Ssz.Capella;
 using Lantern.Beacon.Sync.Types.Ssz.Deneb;
@@ -19,7 +17,9 @@ public class BeaconClient(ISyncProtocol syncProtocol, ILiteDbService liteDbServi
 {
     private readonly ILogger<BeaconClient> _logger = serviceProvider.GetRequiredService<ILoggerFactory>().CreateLogger<BeaconClient>();
 
-    public async Task InitAsync(CancellationToken token = default)
+    public CancellationTokenSource? CancellationTokenSource { get; private set; }
+
+    public async Task InitAsync()
     {
         try
         {
@@ -35,12 +35,12 @@ public async Task InitAsync(CancellationToken token = default)
             peerState.Init(peerFactoryBuilder.AppLayerProtocols);
             gossipSubManager.Init();
 
-            if (gossipSubManager.LightClientFinalityUpdate == null || gossipSubManager.LightClientOptimisticUpdate == null)
+            if (gossipSubManager.LightClientFinalityUpdate == null && gossipSubManager.LightClientOptimisticUpdate == null)
             {
                 return;
             }
 
-            await beaconClientManager.InitAsync(token);
+            await beaconClientManager.InitAsync();
         }
         catch (Exception e)
         {
@@ -53,6 +53,8 @@ public async Task StartAsync(CancellationToken token = default)
     {
         try
         { 
+            CancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(token);
+
             await gossipSubManager.StartAsync(token);
             await beaconClientManager.StartAsync(token); 
         }
@@ -65,8 +67,18 @@ public async Task StartAsync(CancellationToken token = default)
 
     public async Task StopAsync()
     {
+        if (CancellationTokenSource == null)
+        {
+            _logger.LogWarning("Beacon client is not running. Nothing to stop");
+            return;
+        }
+
+        await CancellationTokenSource.CancelAsync();
         await gossipSubManager.StopAsync();
         await beaconClientManager.StopAsync();
+
         liteDbService.Dispose();
+        CancellationTokenSource.Dispose();
+        CancellationTokenSource = null;
     }
 }
diff --git a/src/Lantern.Beacon/BeaconClientManager.cs b/src/Lantern.Beacon/BeaconClientManager.cs
@@ -31,7 +31,7 @@ public class BeaconClientManager(BeaconClientOptions clientOptions,
     public CancellationTokenSource? CancellationTokenSource { get; private set; }
     public ILocalPeer? LocalPeer { get; private set; } 
 
-    public async Task InitAsync(CancellationToken token = default) 
+    public async Task InitAsync() 
     { 
         try 
         { 

diff --git a/src/Lantern.Beacon/BeaconClientPeerFactory.cs b/src/Lantern.Beacon/BeaconClientPeerFactory.cs
@@ -20,10 +20,12 @@ public override ILocalPeer Create(Identity? identity = null, Multiaddress? local
     {
         identity ??= new Identity();
         localAddr ??= $"/ip4/0.0.0.0/tcp/0/p2p/{identity.PeerId}";
+
         if (localAddr.Get<P2P>() is null)
         {
             localAddr.Add<P2P>(identity.PeerId.ToString());
         }
+
         return base.Create(identity, localAddr);
     }
 }
diff --git a/src/Lantern.Beacon/BeaconClientServiceBuilder.cs b/src/Lantern.Beacon/BeaconClientServiceBuilder.cs
@@ -15,20 +15,19 @@ namespace Lantern.Beacon;
 
 public class BeaconClientServiceBuilder(IServiceCollection services) : IBeaconClientServiceBuilder
 {
-    private readonly IDiscv5ProtocolBuilder? _discv5ProtocolBuilder = new Discv5ProtocolBuilder(services);
     private SyncProtocolOptions _syncProtocolOptions = new();
     private BeaconClientOptions _beaconClientOptions = new();
-    private ILoggerFactory _loggerFactory = LoggingOptions.Default;
+    private IDiscv5ProtocolBuilder _discv5ProtocolBuilder = new Discv5ProtocolBuilder(services);
     private IServiceProvider? _serviceProvider;
+    private ILoggerFactory _loggerFactory = LoggingOptions.Default;
 
     public IBeaconClientServiceBuilder AddDiscoveryProtocol(Action<IDiscv5ProtocolBuilder> configure)
     {
         configure(_discv5ProtocolBuilder ?? throw new ArgumentNullException(nameof(configure)));
         return this;
     }
 
-    public IBeaconClientServiceBuilder AddLibp2pProtocol(
-        Func<ILibp2pPeerFactoryBuilder, IPeerFactoryBuilder> factorySetup)
+    public IBeaconClientServiceBuilder AddLibp2pProtocol(Func<ILibp2pPeerFactoryBuilder, IPeerFactoryBuilder> factorySetup)
     {
         services.AddScoped(sp => factorySetup(new BeaconClientPeerFactoryBuilder(sp)))
             .AddScoped<PubsubRouter>()
@@ -75,11 +74,6 @@ public IBeaconClientServiceBuilder WithLoggerFactory(ILoggerFactory loggerFactor
 
     public IBeaconClient Build()
     {
-        if(_discv5ProtocolBuilder == null)
-        {
-            throw new ArgumentNullException(nameof(_discv5ProtocolBuilder));
-        }
-
         services.AddBeaconClient(_discv5ProtocolBuilder.Build(), _beaconClientOptions, _syncProtocolOptions, _loggerFactory);
         _serviceProvider = services.BuildServiceProvider();
 

diff --git a/src/Lantern.Beacon/BeaconClientUtility.cs b/src/Lantern.Beacon/BeaconClientUtility.cs
@@ -56,8 +56,23 @@ private static bool TryGetIpAndPort(IEnr enr, string ipKey, string portKey, out
             return false;
         }
 
-        ip = enr.GetEntry<EntryIp>(ipKey).Value;
-        port = enr.GetEntry<EntryTcp>(portKey).Value;
+        if(ipKey == EnrEntryKey.Ip)
+        {
+            ip = enr.GetEntry<EntryIp>(ipKey).Value;
+        }
+        else if(ipKey == EnrEntryKey.Ip6)
+        {
+            ip = enr.GetEntry<EntryIp6>(ipKey).Value;
+        }
+
+        if(portKey == EnrEntryKey.Tcp)
+        {
+            port = enr.GetEntry<EntryTcp>(portKey).Value;
+        }
+        else if(portKey == EnrEntryKey.Tcp6)
+        {
+            port = enr.GetEntry<EntryTcp6>(portKey).Value;
+        }
 
         return true;
     }

diff --git a/src/Lantern.Beacon/IBeaconClient.cs b/src/Lantern.Beacon/IBeaconClient.cs
@@ -2,7 +2,9 @@ namespace Lantern.Beacon;
 
 public interface IBeaconClient
 {
-    Task InitAsync(CancellationToken token = default);
+    Task InitAsync();
 
     Task StartAsync(CancellationToken token = default);
+
+    Task StopAsync();
 }
diff --git a/src/Lantern.Beacon/IBeaconClientManager.cs b/src/Lantern.Beacon/IBeaconClientManager.cs
@@ -8,7 +8,7 @@ public interface IBeaconClientManager
 
     ILocalPeer? LocalPeer { get; }
 
-    Task InitAsync(CancellationToken token = default);
+    Task InitAsync();
 
     Task StartAsync(CancellationToken token = default);
 

diff --git a/test/Lantern.Beacon.Tests/BeaconChainUtilityTests.cs b/test/Lantern.Beacon.Tests/BeaconChainUtilityTests.cs
@@ -1,10 +1,12 @@
+using System.Net;
 using Lantern.Beacon;
 using Lantern.Beacon.Sync;
 using Lantern.Beacon.Sync.Config;
 using Lantern.Beacon.Sync.Presets;
 using Lantern.Discv5.Enr;
 using Lantern.Discv5.Enr.Entries;
 using Lantern.Discv5.Enr.Identity.V4;
+using Lantern.Discv5.WireProtocol.Session;
 using Multiformats.Address.Protocols;
 using NUnit.Framework;
 using SszSharp;
@@ -15,18 +17,43 @@ namespace Lantern.Beacon.Tests;
 public class BeaconChainUtilityTests
 {
     [Test]
-    public void ConvertToMultiAddress_ShouldCorrectlyConvertToMultiAddress()
+    public void ConvertToMultiAddress_ShouldCorrectlyConvertToMultiAddressForIpV4()
     {
         var enrRegistry = new EnrEntryRegistry();
         var enrString = "enr:-Mq4QLyFLj2R0kwCmxNgO02F2JqHOUAT9CnqK9qHBwJWPlvNR36e9YydkUzFM69E0dzX7hrpOUAJVKsBLb3PysSz-IiGAY7D6Sg4h2F0dG5ldHOIAAAAAAAAAAaEZXRoMpBqlaGpBAAAAP__________gmlkgnY0gmlwhCJkw5SJc2VjcDI1NmsxoQMc6eWKtIsR4Ref474zOEeRKEuHzxrK_jffZrkzzYSuUYhzeW5jbmV0cwCDdGNwgjLIg3VkcILLBIR1ZHA2gi7g";
-
         var enr = new EnrFactory(enrRegistry).CreateFromString(enrString, new IdentityVerifierV4());
         var multiAddress = BeaconClientUtility.ConvertToMultiAddress(enr);
 
         Assert.That(multiAddress, Is.Not.Null);
+        Assert.That(multiAddress.ToString().Contains("ip4"), Is.True);
         Assert.That(enr.GetEntry<EntryIp>(EnrEntryKey.Ip).Value, Is.EqualTo(multiAddress!.Get<IP>().Value));
         Assert.That(enr.GetEntry<EntryTcp>(EnrEntryKey.Tcp).Value, Is.EqualTo(multiAddress.Get<TCP>().Value));
     }
+
+    [Test]
+    public void ConvertToMultiAddress_ShouldCorrectlyConvertToMultiAddressForIpV6()
+    {
+        var enrBuilder = new EnrBuilder();
+        var sessionOptions = SessionOptions.Default;
+
+        enrBuilder.WithIdentityScheme(sessionOptions.Verifier, sessionOptions.Signer);
+        enrBuilder.WithEntry(EnrEntryKey.Id, new EntryId("v4"));
+        enrBuilder.WithEntry(EnrEntryKey.Secp256K1, new EntrySecp256K1(sessionOptions.Signer.PublicKey));
+        enrBuilder.WithEntry(EnrEntryKey.Ip6, new EntryIp6(IPAddress.IPv6Any));
+        enrBuilder.WithEntry(EnrEntryKey.Tcp6, new EntryTcp6(30303));
+
+        var enr = enrBuilder.Build();
+        var multiAddress = BeaconClientUtility.ConvertToMultiAddress(enr);
+
+        Assert.That(multiAddress.ToString().Contains("ip6"), Is.True);
+    }
+
+    [Test]
+    public void ConvertToMultiAddress_ShouldReturnIfEnrIsNull()
+    {
+        var multiAddress = BeaconClientUtility.ConvertToMultiAddress(null);
+        Assert.That(multiAddress, Is.Null);
+    }
 
     [Test]
     public void GetForkDigestBytes_ShouldReturnCorrectForkDigest()

diff --git a/test/Lantern.Beacon.Tests/BeaconClientManagerTests.cs b/test/Lantern.Beacon.Tests/BeaconClientManagerTests.cs
@@ -680,7 +680,7 @@ public async Task RunSyncProtocol_ShouldRunBootstrapProtocolIfSyncProtocolIsNotI
         var peersToDialField = typeof(BeaconClientManager).GetField("_peersToDial", BindingFlags.NonPublic | BindingFlags.Instance);
         var peersToDialQueue = (ConcurrentQueue<Multiaddress>)peersToDialField.GetValue(_beaconClientManager);
 
-        await _beaconClientManager.InitAsync(cts.Token);
+        await _beaconClientManager.InitAsync();
 
         _beaconClientManager.StartAsync(cts.Token);
 
@@ -743,7 +743,7 @@ public async Task RunSyncProtocol_ShouldDisconnectFromPeerIfSyncProtocolDidNotIn
         var peersToDialField = typeof(BeaconClientManager).GetField("_peersToDial", BindingFlags.NonPublic | BindingFlags.Instance);
         var peersToDialQueue = (ConcurrentQueue<Multiaddress>)peersToDialField.GetValue(_beaconClientManager);
 
-        await _beaconClientManager.InitAsync(cts.Token);
+        await _beaconClientManager.InitAsync();
 
         _beaconClientManager.StartAsync(cts.Token);
 
@@ -808,7 +808,7 @@ public async Task RunSyncProtocol_ShouldSyncDenebForkIfActiveForkIsSetToDeneb()
         var peersToDialField = typeof(BeaconClientManager).GetField("_peersToDial", BindingFlags.NonPublic | BindingFlags.Instance);
         var peersToDialQueue = (ConcurrentQueue<Multiaddress>)peersToDialField.GetValue(_beaconClientManager);
 
-        await _beaconClientManager.InitAsync(cts.Token);
+        await _beaconClientManager.InitAsync();
         _beaconClientManager.StartAsync(cts.Token);
         await Task.Delay(1000, cts.Token);
 
@@ -870,7 +870,7 @@ public async Task RunSyncProtocol_ShouldSyncCapellaForkIfActiveForkIsSetToCapell
         var peersToDialField = typeof(BeaconClientManager).GetField("_peersToDial", BindingFlags.NonPublic | BindingFlags.Instance);
         var peersToDialQueue = (ConcurrentQueue<Multiaddress>)peersToDialField.GetValue(_beaconClientManager);
 
-        await _beaconClientManager.InitAsync(cts.Token);
+        await _beaconClientManager.InitAsync();
         _beaconClientManager.StartAsync(cts.Token);
         await Task.Delay(1000, cts.Token);
 
@@ -932,7 +932,7 @@ public async Task RunSyncProtocol_ShouldBellatrixForkIfActiveForkIsSetToBellatri
         var peersToDialField = typeof(BeaconClientManager).GetField("_peersToDial", BindingFlags.NonPublic | BindingFlags.Instance);
         var peersToDialQueue = (ConcurrentQueue<Multiaddress>)peersToDialField.GetValue(_beaconClientManager);
 
-        await _beaconClientManager.InitAsync(cts.Token);
+        await _beaconClientManager.InitAsync();
         _beaconClientManager.StartAsync(cts.Token);
         await Task.Delay(1000, cts.Token);
 
@@ -994,7 +994,7 @@ public async Task RunSyncProtocol_ShouldAltairForkIfActiveForkIsSetToAltair()
         var peersToDialField = typeof(BeaconClientManager).GetField("_peersToDial", BindingFlags.NonPublic | BindingFlags.Instance);
         var peersToDialQueue = (ConcurrentQueue<Multiaddress>)peersToDialField.GetValue(_beaconClientManager);
 
-        await _beaconClientManager.InitAsync(cts.Token);
+        await _beaconClientManager.InitAsync();
         _beaconClientManager.StartAsync(cts.Token);
         await Task.Delay(1000, cts.Token);
 

diff --git a/test/Lantern.Beacon.Tests/BeaconClientPeerFactoryTests.cs b/test/Lantern.Beacon.Tests/BeaconClientPeerFactoryTests.cs
@@ -0,0 +1,46 @@
+using System.Reflection;
+using Lantern.Beacon.Networking.Libp2pProtocols.Identify;
+using Moq;
+using Nethermind.Libp2p.Core;
+using NUnit.Framework;
+
+namespace Lantern.Beacon.Tests;
+
+[TestFixture]
+public class BeaconClientPeerFactoryTests
+{
+    private BeaconClientPeerFactory _beaconClientPeerFactory;
+
+    [SetUp]
+    public void Setup()
+    {
+        _beaconClientPeerFactory = new BeaconClientPeerFactory(null);
+    }
+
+    [Test]
+    public async Task ConnectedTo_ShouldDialPeerIdentifyProtocol()
+    {
+        var mockRemotePeer = new Mock<IRemotePeer>();
+        var method = typeof(BeaconClientPeerFactory).GetMethod("ConnectedTo", BindingFlags.NonPublic | BindingFlags.Instance);
+
+        if (method == null)
+        {
+            Assert.Fail("The method 'ConnectedTo' was not found.");
+        }
+        else
+        {
+            var task = (Task)method.Invoke(_beaconClientPeerFactory, [mockRemotePeer.Object, true]);
+            await task;
+        }
+
+        mockRemotePeer.Verify(x => x.DialAsync<PeerIdentifyProtocol>(new CancellationToken()), Times.Once());
+    }
+
+    [Test]
+    public void Create_ShouldReturnLocalPeer()
+    {
+        var localPeer = _beaconClientPeerFactory.Create();
+
+        Assert.That(localPeer, Is.Not.Null);
+    }
+}