Add feature to ignore only some unhandled messages

JKorf · Mar 22, 2024 · 1bc2251 · 1bc2251
1 parent 108c8fc
commit 1bc2251
Show file tree

Hide file tree

Showing 5 changed files with 60 additions and 3 deletions.
diff --git a/CryptoExchange.Net.UnitTests/SocketClientTests.cs b/CryptoExchange.Net.UnitTests/SocketClientTests.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using CryptoExchange.Net.Objects;
@@ -233,5 +234,31 @@ public async Task SuccessResponse_Should_ConfirmSubscription()
             // assert
             Assert.That(client.SubClient.TestSubscription.Confirmed);
         }
+
+        [TestCase()]
+        public void ExpectedUnhandledMessaged_Should_BeIgnored()
+        {
+            // arrange
+            var client = new TestSocketClient(opt =>
+            {
+                opt.OutputOriginalData = true;
+            });
+            var socket = client.CreateSocket();
+            socket.CanConnect = true;
+            var connection = new SocketConnection(new TraceLogger(), client.SubClient, socket, "https://test.test");
+            client.SubClient.ConnectSocketSub(connection);
+            var subObj = new TestSubscription<Dictionary<string, string>>(Mock.Of<ILogger>(), messageEvent => { });
+            connection.AddSubscription(subObj);
+
+            // act
+            client.SubClient.IgnoredUnhandledMessages = new HashSet<string> { "connected" };
+            var unhandledMessage = JsonConvert.SerializeObject(new { topic = "unhandled" });
+            socket.InvokeMessage(JsonConvert.SerializeObject(new { topic = "connected" }));
+            socket.InvokeMessage(unhandledMessage);
+
+            // assert
+            Assert.That(client.SubClient.UnhandledMessages.Count == 1);
+            Assert.That(client.SubClient.UnhandledMessages.First().Equals(unhandledMessage));
+        }
     }
 }
diff --git a/CryptoExchange.Net.UnitTests/TestImplementations/TestSocketClient.cs b/CryptoExchange.Net.UnitTests/TestImplementations/TestSocketClient.cs
@@ -13,6 +13,7 @@
 using CryptoExchange.Net.UnitTests.TestImplementations.Sockets;
 using Microsoft.Extensions.Logging;
 using Moq;
+using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
 
 namespace CryptoExchange.Net.UnitTests.TestImplementations
@@ -74,6 +75,7 @@ public class TestSocketOptions: SocketExchangeOptions<TestEnvironment>
 
     public class TestSubSocketClient : SocketApiClient
     {
+        public List<object> UnhandledMessages { get; } = new List<object>();
         private MessagePath _channelPath = MessagePath.Get().Property("channel");
         private MessagePath _topicPath = MessagePath.Get().Property("topic");
 
@@ -84,6 +86,18 @@ public TestSubSocketClient(TestSocketOptions options, SocketApiOptions apiOption
 
         }
 
+        protected override void HandleUnhandledMessage(IMessageAccessor message)
+        {
+            if (message.Underlying is JObject json)
+            {
+                UnhandledMessages.Add(json.ToString(Formatting.None));
+            }
+            else
+            {
+                UnhandledMessages.Add(message.Underlying);
+            }
+        }
+
         internal IWebsocket CreateSocketInternal(string address)
         {
             return CreateSocket(address);

diff --git a/CryptoExchange.Net/Clients/SocketApiClient.cs b/CryptoExchange.Net/Clients/SocketApiClient.cs
@@ -53,6 +53,11 @@ public abstract class SocketApiClient : BaseApiClient, ISocketApiClient
         /// </summary>
         protected internal bool UnhandledMessageExpected { get; set; }
 
+        /// <summary>
+        /// To be used if <see cref="UnhandledMessageExpected"/> is false and some messages should be ignored
+        /// </summary>
+        protected internal HashSet<string>? IgnoredUnhandledMessages { get; set; } = null;
+
         /// <summary>
         /// If true a subscription will accept message before the confirmation of a subscription has been received
         /// </summary>
@@ -496,6 +501,16 @@ protected virtual void HandleUnhandledMessage(IMessageAccessor message)
         {
         }
 
+        /// <summary>
+        /// Check if an unhandled message with the given listen id should be ignored
+        /// </summary>
+        /// <param name="listenId"></param>
+        /// <returns></returns>
+        public bool IsUnhandledMessageIgnored(string listenId)
+        {
+            return IgnoredUnhandledMessages?.Contains(listenId) ?? false;
+        }
+
         /// <summary>
         /// Connect a socket
         /// </summary>
@@ -505,6 +520,7 @@ protected virtual async Task<CallResult<bool>> ConnectSocketAsync(SocketConnecti
         {
             if (await socketConnection.ConnectAsync().ConfigureAwait(false))
             {
+                socketConnection.UnhandledMessage += HandleUnhandledMessage;
                 socketConnections.TryAdd(socketConnection.SocketId, socketConnection);
                 return new CallResult<bool>(true);
             }

diff --git a/CryptoExchange.Net/Objects/TraceLogger.cs b/CryptoExchange.Net/Objects/TraceLogger.cs
@@ -49,11 +49,11 @@ public TraceLogger(string? categoryName = null, LogLevel level = LogLevel.Trace)
         public IDisposable? BeginScope<TState>(TState state) where TState : notnull => null!;
 
         /// <inheritdoc />
-        public bool IsEnabled(LogLevel logLevel) => (int)logLevel < (int)_logLevel;
+        public bool IsEnabled(LogLevel logLevel) => logLevel.CompareTo(_logLevel) >= 0;
         /// <inheritdoc />
         public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
         {
-            if ((int)logLevel < (int)_logLevel)
+            if (!IsEnabled(_logLevel))
                 return;
 
             var logMessage = $"{DateTime.Now:yyyy/MM/dd HH:mm:ss:fff} | {logLevel} | {(_categoryName == null ? "" : $"{_categoryName} | ")}{formatter(state, exception)}";

diff --git a/CryptoExchange.Net/Sockets/SocketConnection.cs b/CryptoExchange.Net/Sockets/SocketConnection.cs
@@ -403,7 +403,7 @@ protected virtual void HandleStreamMessage(WebSocketMessageType type, ReadOnlyMe
 
                 if (processors.Count == 0)
                 {
-                    if (!ApiClient.UnhandledMessageExpected)
+                    if (!ApiClient.UnhandledMessageExpected && !ApiClient.IsUnhandledMessageIgnored(listenId))
                     {
                         List<string> listenerIds;
                         lock (_listenersLock)