如何对含多后台持续任务的WpaConnection类进行单元测试

问题描述 #

需要对包含多个持续后台任务的WpaConnection类开展单元测试，该类中任务的启动依赖TCP连接。希望了解针对此类的最佳单元测试方案，以及如何Mock TCP Client以跳过ConnectAsync调用。

原代码示例 #

using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace WpaConnection
{
        public class WpaConnection : IDisposable
        {
            private readonly IPEndPoint _wapIpEndPoint;
            private readonly TcpClient _tcpClient = new TcpClient();
            private readonly Task _connectTask;
            private bool _isExiting;

            public WpaConnection(
                IPEndPoint wapIpEndPoint,
                CancellationToken stoppingToken)
            {
                _wapIpEndPoint = wapIpEndPoint;
                _connectTask = ConnectTask(wapIpEndPoint, stoppingToken);
                _isExiting = false;
            }

            public event EventHandler<WapReadEventArgs> OnRead;

            public void Dispose()
            {
                Dispose(true);
                GC.SuppressFinalize(this);
            }

            public void Disconnect()
            {
                if (!_connectTask.IsCompleted)
                {
                    _isExiting = true;
                }
            }

            protected virtual void Dispose(bool disposing)
            {
                if (disposing)
                {
                    _tcpClient.Dispose();
                }
            }

            private async Task ConnectTask(IPEndPoint wapIpEndPoint, CancellationToken stoppingToken)
            {
                CancellationTokenSource cancellationTokenSource = new CancellationTokenSource();
                Task readEventAsyncTask;
                Task task1;
                Task task2;
                Task task3;

                Console.WriteLine("WpaConnection: ConnectTask started");

                try
                {
                    await _tcpClient.ConnectAsync(wapIpEndPoint.Address, wapIpEndPoint.Port);
                    Console.WriteLine("WpaConnection: ConnectTask Connected to {WapEndPoint}", wapIpEndPoint);


                    readEventAsyncTask = ReadEventTask(cancellationTokenSource.Token);
                    task1 = Task1(cancellationTokenSource.Token);
                    task2 = Task2(cancellationTokenSource.Token);
                    task3 = Task3(cancellationTokenSource.Token);

                    while (!stoppingToken.IsCancellationRequested && _tcpClient.Connected && !_isExiting)
                    {
                        await Task.Delay(1, stoppingToken);
                    }

                    _tcpClient.Close();

                    Console.WriteLine("WpaConnection: ConnectTask closed TcpClient");
                }
                catch (Exception)
                {
                    _isExiting = true;
                }

                cancellationTokenSource.Cancel();


                Console.WriteLine("WpaConnection: ConnectTask exiting...");
            }

            private async Task Task3(CancellationToken stoppingToken)
            {
                Console.WriteLine("WpaConnection: Task3 started");

                try
                {
                    int sequence = 0;
                    while (!stoppingToken.IsCancellationRequested && !_isExiting)
                    {
                        Console.WriteLine($"Task3 {sequence++}");

                        await Task.Delay(1000, stoppingToken);
                    }
                }
                catch (Exception)
                {
                    _isExiting = true;
                }

                Console.WriteLine("WpaConnection: Task3 exiting...");
            }

            private async Task Task1(CancellationToken stoppingToken)
            {
                Console.WriteLine("WpaConnection: Task1 started");

                try
                {
                    int sequence = 0;
                    while (!stoppingToken.IsCancellationRequested && !_isExiting)
                    {
                        Console.WriteLine($"Task1 {sequence++}");

                        await Task.Delay(1000, stoppingToken);
                    }
                }
                catch (Exception)
                {
                    _isExiting = true;
                }

                Console.WriteLine("WpaConnection: Task1 exiting...");
            }

            private async Task Task2(CancellationToken stoppingToken)
            {
                Console.WriteLine("WpaConnection: Task2 started");

                try
                {
                    int sequence = 0;
                    while (!stoppingToken.IsCancellationRequested && !_isExiting)
                    {
                        Console.WriteLine($"Task2 {sequence++}");

                        await Task.Delay(1000, stoppingToken);
                    }
                }
                catch (Exception)
                {
                    _isExiting = true;
                }

                Console.WriteLine("WpaConnection: Task2 exiting...");
            }

            private async Task ReadEventTask(CancellationToken stoppingToken)
            {
                Console.WriteLine("WpaConnection: ReadEventTask started");
                byte[] buffer = new byte[512];

                try
                {
                    while (!stoppingToken.IsCancellationRequested && !_isExiting && _tcpClient.Connected)
                    {
                        int length = await _tcpClient.GetStream().ReadAsync(buffer, 0, buffer.Length, stoppingToken);
                        if (length > 0)
                        {
                            byte[] data = new byte[length];
                            Buffer.BlockCopy(buffer, 0, data, 0, length);

                            OnRead?.Invoke(
                                this,
                                new WapReadEventArgs()
                                {
                                    WapEndPoint = _wapIpEndPoint,
                                    Length = length,
                                    Buffer = data,
                                });
                        }
                    }
                }
                catch (Exception)
                {
                    _isExiting = true;
                }

                Console.WriteLine("WpaConnection: ReadEventTask exiting...");
            }
        }

        public class WapReadEventArgs : EventArgs
        {
            public IPEndPoint WapEndPoint { get; set; }

            public int Length { get; set; }

            public byte[] Buffer { get; set; }
        }

    internal class Program
    {
        static void Main(string[] args)
        {
            IPEndPoint iPEndPoint = new IPEndPoint(IPAddress.Any, 1234);
            CancellationTokenSource cancellationToken = new CancellationTokenSource();
            WpaConnection connection = new WpaConnection(iPEndPoint, cancellationToken.Token);

            Task.Delay(1000).Wait();

            cancellationToken.Cancel();

            Console.ReadLine();
        }
    }
}

解决方案 #

1. 重构代码实现依赖注入 #

原代码中TcpClient是硬实例化的，无法直接Mock。首先抽象TCP客户端的行为，定义接口并包装原生TcpClient：

namespace WpaConnection
{
    public interface ITcpClient : IDisposable
    {
        Task ConnectAsync(IPAddress address, int port, CancellationToken cancellationToken = default);
        NetworkStream GetStream();
        bool Connected { get; }
        void Close();
    }

    public class TcpClientWrapper : ITcpClient
    {
        private readonly TcpClient _tcpClient;

        public TcpClientWrapper() => _tcpClient = new TcpClient();
        public bool Connected => _tcpClient.Connected;

        public async Task ConnectAsync(IPAddress address, int port, CancellationToken cancellationToken = default)
        {
            await _tcpClient.ConnectAsync(address, port, cancellationToken);
        }

        public void Close() => _tcpClient.Close();

        public void Dispose() => _tcpClient.Dispose();

        public NetworkStream GetStream() => _tcpClient.GetStream();
    }
}

修改WpaConnection的构造函数，注入ITcpClient替换硬实例化的TcpClient：

public class WpaConnection : IDisposable
{
    private readonly IPEndPoint _wapIpEndPoint;
    private readonly ITcpClient _tcpClient;
    private readonly Task _connectTask;
    private bool _isExiting;

    public WpaConnection(
        IPEndPoint wapIpEndPoint,
        ITcpClient tcpClient,
        CancellationToken stoppingToken)
    {
        _wapIpEndPoint = wapIpEndPoint;
        _tcpClient = tcpClient;
        _connectTask = ConnectTask(wapIpEndPoint, stoppingToken);
        _isExiting = false;
    }

    // 保留原有其他方法，仅替换所有_tcpClient的引用为接口实例
    // ... 其余代码不变 ...
}

2. 使用Mock框架模拟TCP连接 #

以Moq框架为例，模拟ITcpClient的行为，跳过真实TCP连接并模拟数据流：

using Moq;
using Xunit;
using System.Net.Sockets;
using System.Threading.Tasks;
using System.Threading;
using System;

namespace WpaConnection.Tests
{
    public class WpaConnectionTests
    {
        [Fact]
        public async Task ConnectTask_ShouldStartBackgroundTasks_WhenConnected()
        {
            // 准备Mock对象
            var mockTcpClient = new Mock<ITcpClient>();
            var mockStream = new Mock<NetworkStream>(MockBehavior.Loose, new System.IO.MemoryStream());
            
            // 模拟ConnectAsync立即成功
            mockTcpClient.Setup(c => c.ConnectAsync(It.IsAny<IPAddress>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
                         .Returns(Task.CompletedTask);
            mockTcpClient.Setup(c => c.Connected).Returns(true);
            mockTcpClient.Setup(c => c.GetStream()).Returns(mockStream.Object);

            var endPoint = new IPEndPoint(IPAddress.Loopback, 1234);
            using var cts = new CancellationTokenSource();
            var connection = new WpaConnection(endPoint, mockTcpClient.Object, cts.Token);

            // 等待后台任务启动
            await Task.Delay(100);

            // 此处可根据实际逻辑添加断言，例如验证任务是否启动（可通过日志或内部状态判断）
            Assert.True(true);

            // 清理资源
            cts.Cancel();
            connection.Dispose();
        }

        [Fact]
        public async Task ReadEventTask_ShouldTriggerOnRead_WhenDataReceived()
        {
            var mockTcpClient = new Mock<ITcpClient>();
            var testData = Encoding.UTF8.GetBytes("test data");
            var memoryStream = new System.IO.MemoryStream(testData);
            var mockStream = new NetworkStream(memoryStream);

            mockTcpClient.Setup(c => c.ConnectAsync(It.IsAny<IPAddress>(), It.IsAny<int>(), It.IsAny<CancellationToken>()))
                         .Returns(Task.CompletedTask);
            mockTcpClient.Setup(c => c.Connected).Returns(true);
            mockTcpClient.Setup(c => c.GetStream()).Returns(mockStream);

            var endPoint = new IPEndPoint(IPAddress.Loopback, 1234);
            using var cts = new CancellationTokenSource();
            var connection = new WpaConnection(endPoint, mockTcpClient.Object, cts.Token);

            WapReadEventArgs receivedArgs = null;
            connection.OnRead += (s, e) => receivedArgs = e;

            // 等待数据读取触发事件
            await Task.Delay(100);

            Assert.NotNull(receivedArgs);
            Assert.Equal(testData.Length, receivedArgs.Length);
            Assert.Equal(testData, receivedArgs.Buffer);

            cts.Cancel();
            connection.Dispose();
        }
    }
}

3. 单元测试核心场景 #

• 连接成功场景：验证后台任务（Task1/Task2/Task3/ReadEventTask）是否正确启动
• 数据接收场景：模拟数据流，验证OnRead事件是否正确触发并传递数据
• 断开连接场景：调用Disconnect或触发取消令牌，验证后台任务是否正确终止、TCP连接是否关闭
• 异常场景：模拟连接失败、数据流异常，验证_isExiting状态是否正确设置，任务是否优雅退出

内容的提问来源于stack exchange，提问作者Hamza