FlowScope/My project/Assets/FlowScope/Tests/PlayMode/UI/UITestHelpers.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using FlowScope.Resources;
using NUnit.Framework;
using UnityEngine;

namespace FlowScope.Tests.PlayMode.UI
{
    internal sealed class UITestResources : IResourceService
    {
        private readonly Dictionary<string, GameObject> _assets = new();

        public List<string> LoadedKeys { get; } = new();
        public List<UITestResourceHandle<GameObject>> Handles { get; } = new();

        public void Register(string key, GameObject asset)
        {
            _assets.Add(key, asset);
        }

        public Task<IResourceHandle<T>> LoadAsync<T>(string key, CancellationToken cancellationToken)
            where T : class
        {
            cancellationToken.ThrowIfCancellationRequested();
            LoadedKeys.Add(key);

            if (!_assets.TryGetValue(key, out var asset))
            {
                throw new InvalidOperationException($"Missing resource: {key}");
            }

            var handle = new UITestResourceHandle<GameObject>(key, asset);
            Handles.Add(handle);
            return Task.FromResult<IResourceHandle<T>>((IResourceHandle<T>)(object)handle);
        }

        public IResourceGroup CreateGroup()
        {
            throw new NotSupportedException();
        }
    }

    internal sealed class UITestResourceHandle<T> : IResourceHandle<T>
        where T : class
    {
        public UITestResourceHandle(string key, T asset)
        {
            Key = key;
            Asset = asset;
        }

        public string Key { get; }
        public T Asset { get; private set; }
        public bool IsDisposed { get; private set; }

        public void Dispose()
        {
            IsDisposed = true;
        }
    }

    internal static class UITestAsync
    {
        public static IEnumerator Await(Task task)
        {
            while (!task.IsCompleted)
            {
                yield return null;
            }

            if (task.IsFaulted)
            {
                throw task.Exception.GetBaseException();
            }

            if (task.IsCanceled)
            {
                throw new OperationCanceledException();
            }
        }

        public static IEnumerator Await<T>(Task<T> task, Action<T> onCompleted)
        {
            while (!task.IsCompleted)
            {
                yield return null;
            }

            if (task.IsFaulted)
            {
                throw task.Exception.GetBaseException();
            }

            if (task.IsCanceled)
            {
                throw new OperationCanceledException();
            }

            onCompleted(task.Result);
        }

        public static IEnumerator AwaitExpected<TException>(Task task)
            where TException : Exception
        {
            while (!task.IsCompleted)
            {
                yield return null;
            }

            Assert.That(task.IsFaulted || task.IsCanceled, Is.True);
            if (task.IsCanceled)
            {
                Assert.That(typeof(TException), Is.EqualTo(typeof(OperationCanceledException)));
                yield break;
            }

            Assert.That(task.Exception.GetBaseException(), Is.TypeOf<TException>());
        }
    }
}