Axi_Omron/tests/AxiOmron.PcbCheck.Tests/ShutdownHelperTests.cs

using System.Collections.Concurrent;
using AxiOmron.PcbCheck.Utils;
using Microsoft.Extensions.Hosting;
using Xunit;

namespace AxiOmron.PcbCheck.Tests;

/// <summary>
/// 验证应用退出阶段的异步清理流程不会丢失原始线程上下文。
/// </summary>
public sealed class ShutdownHelperTests
{
    /// <summary>
    /// 停止 Host 后，后续收尾回调应继续在捕获到的同步上下文线程上执行。
    /// </summary>
    /// <returns>异步测试任务。</returns>
    [Fact]
    public async Task StopHostAsync_ShouldResumeOnCapturedSynchronizationContext()
    {
        await SingleThreadSynchronizationContext.RunAsync(async () =>
        {
            int expectedThreadId = Environment.CurrentManagedThreadId;
            int callbackThreadId = -1;
            var host = new FakeHost();

            await ShutdownHelper.StopHostAsync(
                host,
                TimeSpan.FromMilliseconds(200),
                () => callbackThreadId = Environment.CurrentManagedThreadId);

            Assert.True(host.StopAsyncCalled);
            Assert.Equal(expectedThreadId, callbackThreadId);
        });
    }

    private sealed class FakeHost : IHost
    {
        public IServiceProvider Services => throw new NotSupportedException();

        public bool StopAsyncCalled { get; private set; }

        public void Dispose()
        {
        }

        public Task StartAsync(CancellationToken cancellationToken = default)
        {
            return Task.CompletedTask;
        }

        public async Task StopAsync(CancellationToken cancellationToken = default)
        {
            StopAsyncCalled = true;
            await Task.Yield();
        }
    }

    /// <summary>
    /// 提供仅在单线程内串行泵送回调的同步上下文，便于验证 await 后的线程恢复行为。
    /// </summary>
    private sealed class SingleThreadSynchronizationContext : SynchronizationContext
    {
        private readonly BlockingCollection<(SendOrPostCallback Callback, object? State)> _queue = new();

        /// <summary>
        /// 在专用线程同步上下文中执行异步委托。
        /// </summary>
        /// <param name="action">待执行的异步委托。</param>
        /// <returns>表示执行完成的任务。</returns>
        public static Task RunAsync(Func<Task> action)
        {
            ArgumentNullException.ThrowIfNull(action);

            var completionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
            var thread = new Thread(() =>
            {
                var context = new SingleThreadSynchronizationContext();
                SynchronizationContext.SetSynchronizationContext(context);
                Task task;

                try
                {
                    task = action();
                }
                catch (Exception ex)
                {
                    completionSource.SetException(ex);
                    return;
                }

                task.ContinueWith(
                    completedTask =>
                    {
                        if (completedTask.IsFaulted)
                        {
                            completionSource.SetException(completedTask.Exception!.InnerExceptions);
                        }
                        else if (completedTask.IsCanceled)
                        {
                            completionSource.SetCanceled();
                        }
                        else
                        {
                            completionSource.SetResult();
                        }

                        context._queue.CompleteAdding();
                    },
                    CancellationToken.None,
                    TaskContinuationOptions.None,
                    TaskScheduler.Default);

                foreach ((SendOrPostCallback callback, object? state) in context._queue.GetConsumingEnumerable())
                {
                    callback(state);
                }
            })
            {
                IsBackground = true
            };

            thread.SetApartmentState(ApartmentState.STA);
            thread.Start();
            return completionSource.Task;
        }

        /// <summary>
        /// 将委托异步排入当前上下文队列。
        /// </summary>
        /// <param name="d">待执行回调。</param>
        /// <param name="state">状态对象。</param>
        public override void Post(SendOrPostCallback d, object? state)
        {
            ArgumentNullException.ThrowIfNull(d);
            _queue.Add((d, state));
        }

        /// <summary>
        /// 在当前线程直接执行同步回调。
        /// </summary>
        /// <param name="d">待执行回调。</param>
        /// <param name="state">状态对象。</param>
        public override void Send(SendOrPostCallback d, object? state)
        {
            ArgumentNullException.ThrowIfNull(d);
            d(state);
        }
    }
}