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

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

namespace FlowScope.Tests.PlayMode.UI
{
    public sealed class UIPreloadServiceTests
    {
        private readonly List<GameObject> _objects = new();

        [TearDown]
        public void TearDown()
        {
            foreach (var obj in _objects)
            {
                if (obj != null)
                {
                    UnityEngine.Object.DestroyImmediate(obj);
                }
            }

            _objects.Clear();
        }

        [UnityTest]
        public IEnumerator PreloadAsync_LoadsPanelResourceAndKeepsHandle()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new UITestResources();
            resources.Register("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            yield return UITestAsync.Await(preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None));

            Assert.That(resources.LoadedKeys, Is.EqualTo(new[] { "Custom/Explicit" }));
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.True);
            Assert.That(preload.TryGetPreloadedHandle<ExplicitPathPanel>(out var handle), Is.True);
            Assert.That(handle.Asset, Is.SameAs(prefab));
            Assert.That(resources.Handles[0].IsDisposed, Is.False);
        }

        [UnityTest]
        public IEnumerator PreloadAsync_WhenRepeated_DoesNotLoadAgain()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new UITestResources();
            resources.Register("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            yield return UITestAsync.Await(preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None));
            yield return UITestAsync.Await(preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None));

            Assert.That(resources.LoadedKeys, Is.EqualTo(new[] { "Custom/Explicit" }));
            Assert.That(resources.Handles, Has.Count.EqualTo(1));
        }

        [UnityTest]
        public IEnumerator PreloadAsync_WhenSamePanelLoadsConcurrently_UsesOneResourceLoad()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            var first = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            var second = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);

            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            resources.Complete();
            yield return UITestAsync.Await(first);
            yield return UITestAsync.Await(second);

            Assert.That(resources.LoadCalls, Is.EqualTo(1));
            Assert.That(resources.Handle.IsDisposed, Is.False);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.True);
        }

        [UnityTest]
        public IEnumerator PreloadAsync_WhenFirstWaiterCancels_SecondWaiterStillSucceeds()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);
            using var firstCancellation = new CancellationTokenSource();

            var first = preload.PreloadAsync<ExplicitPathPanel>(firstCancellation.Token);
            var second = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);

            Assert.That(resources.LoadCalls, Is.EqualTo(1));
            Assert.That(resources.LoadCancellationTokenCanBeCanceled, Is.False);

            firstCancellation.Cancel();
            yield return UITestAsync.AwaitExpected<OperationCanceledException>(first);

            resources.Complete();
            yield return UITestAsync.Await(second);

            Assert.That(resources.LoadCalls, Is.EqualTo(1));
            Assert.That(resources.Handle.IsDisposed, Is.False);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.True);
        }

        [UnityTest]
        public IEnumerator PreloadAsync_WhenSecondWaiterCancels_FirstWaiterStillSucceeds()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);
            using var secondCancellation = new CancellationTokenSource();

            var first = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            var second = preload.PreloadAsync<ExplicitPathPanel>(secondCancellation.Token);

            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            secondCancellation.Cancel();
            yield return UITestAsync.AwaitExpected<OperationCanceledException>(second);

            resources.Complete();
            yield return UITestAsync.Await(first);

            Assert.That(resources.LoadCalls, Is.EqualTo(1));
            Assert.That(resources.Handle.IsDisposed, Is.False);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.True);
        }

        [UnityTest]
        public IEnumerator PreloadAsync_WhenAllWaitersCancel_ReleasesLoadedHandleWithoutCaching()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);
            using var firstCancellation = new CancellationTokenSource();
            using var secondCancellation = new CancellationTokenSource();

            var first = preload.PreloadAsync<ExplicitPathPanel>(firstCancellation.Token);
            var second = preload.PreloadAsync<ExplicitPathPanel>(secondCancellation.Token);

            firstCancellation.Cancel();
            secondCancellation.Cancel();
            yield return UITestAsync.AwaitExpected<OperationCanceledException>(first);
            yield return UITestAsync.AwaitExpected<OperationCanceledException>(second);

            resources.Complete();
            yield return null;

            Assert.That(resources.LoadCalls, Is.EqualTo(1));
            Assert.That(resources.Handle.IsDisposed, Is.True);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);
        }

        [UnityTest]
        public IEnumerator ReleaseAll_WhenLoadInFlight_DoesNotCacheAndDisposesLoadedHandle()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            var task = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            preload.ReleaseAll();
            resources.Complete();
            yield return UITestAsync.Await(task);

            Assert.That(resources.Handle.IsDisposed, Is.True);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);
        }

        [UnityTest]
        public IEnumerator PreloadAsync_AfterReleaseAllWhileOldLoadInFlight_StartsNewLoadAndCachesNewHandle()
        {
            var oldPrefab = CreatePrefab<ExplicitPathPanel>("OldExplicitPathPrefab");
            var newPrefab = CreatePrefab<ExplicitPathPanel>("NewExplicitPathPrefab");
            var resources = new SequencedGatedUITestResources("Custom/Explicit", oldPrefab, newPrefab);
            var preload = new UIPreloadService(resources);

            var oldTask = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            preload.ReleaseAll();

            var newTask = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(2));

            resources.Complete(0);
            yield return UITestAsync.Await(oldTask);

            Assert.That(resources.HandleAt(0).IsDisposed, Is.True);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);

            resources.Complete(1);
            yield return UITestAsync.Await(newTask);

            Assert.That(resources.HandleAt(1).IsDisposed, Is.False);
            Assert.That(preload.TryGetPreloadedHandle<ExplicitPathPanel>(out var handle), Is.True);
            Assert.That(handle.Asset, Is.SameAs(newPrefab));
        }

        [UnityTest]
        public IEnumerator OldInFlightCompletion_AfterReleaseAll_DoesNotRemoveNewInFlight()
        {
            var oldPrefab = CreatePrefab<ExplicitPathPanel>("OldExplicitPathPrefab");
            var newPrefab = CreatePrefab<ExplicitPathPanel>("NewExplicitPathPrefab");
            var resources = new SequencedGatedUITestResources("Custom/Explicit", oldPrefab, newPrefab);
            var preload = new UIPreloadService(resources);

            var oldTask = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            preload.ReleaseAll();

            var newTask = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(2));

            resources.Complete(0);
            yield return UITestAsync.Await(oldTask);

            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);

            var thirdTask = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(2));

            resources.Complete(1);
            yield return UITestAsync.Await(newTask);
            yield return UITestAsync.Await(thirdTask);

            Assert.That(resources.HandleAt(1).IsDisposed, Is.False);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.True);
        }

        [UnityTest]
        public IEnumerator Dispose_WhenLoadInFlight_DoesNotCacheAndDisposesLoadedHandle()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            var task = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            preload.Dispose();
            resources.Complete();
            yield return UITestAsync.Await(task);

            Assert.That(resources.Handle.IsDisposed, Is.True);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);
        }

        [UnityTest]
        public IEnumerator PreloadAsync_AfterDispose_ThrowsObjectDisposedException()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new UITestResources();
            resources.Register("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            preload.Dispose();

            yield return UITestAsync.AwaitExpected<ObjectDisposedException>(
                preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None));
            Assert.That(resources.LoadedKeys, Is.Empty);
        }

        [UnityTest]
        public IEnumerator OpenAsync_WhenPanelIsPreloaded_ReusesPreloadedHandle()
        {
            var canvas = CreateCanvas("popup");
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new UITestResources();
            resources.Register("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);
            var manager = new UIManager(resources, preload);
            manager.RegisterLayer("popup", canvas, 0);

            yield return UITestAsync.Await(preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None));
            yield return UITestAsync.Await(
                manager.OpenAsync<ExplicitPathPanel, TestViewModel>(
                    new TestViewModel("preloaded"),
                    CancellationToken.None));
            manager.Close<ExplicitPathPanel>();

            Assert.That(resources.LoadedKeys, Is.EqualTo(new[] { "Custom/Explicit" }));
            Assert.That(resources.Handles[0].IsDisposed, Is.False);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.True);
        }

        [UnityTest]
        public IEnumerator Release_DisposesPreloadedHandle()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new UITestResources();
            resources.Register("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            yield return UITestAsync.Await(preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None));
            preload.Release<ExplicitPathPanel>();

            Assert.That(resources.Handles[0].IsDisposed, Is.True);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);
        }

        [UnityTest]
        public IEnumerator Release_WhenLoadInFlight_DoesNotCacheAndDisposesLoadedHandle()
        {
            var prefab = CreatePrefab<ExplicitPathPanel>("ExplicitPathPrefab");
            var resources = new GatedUITestResources("Custom/Explicit", prefab);
            var preload = new UIPreloadService(resources);

            var task = preload.PreloadAsync<ExplicitPathPanel>(CancellationToken.None);
            Assert.That(resources.LoadCalls, Is.EqualTo(1));

            preload.Release<ExplicitPathPanel>();
            resources.Complete();
            yield return UITestAsync.Await(task);

            Assert.That(resources.Handle.IsDisposed, Is.True);
            Assert.That(preload.IsPreloaded<ExplicitPathPanel>(), Is.False);
        }

        private Canvas CreateCanvas(string name)
        {
            var gameObject = new GameObject(name);
            _objects.Add(gameObject);
            return gameObject.AddComponent<Canvas>();
        }

        private GameObject CreatePrefab<TPanel>(string name)
            where TPanel : Component
        {
            var gameObject = new GameObject(name);
            _objects.Add(gameObject);
            gameObject.AddComponent<TPanel>();
            return gameObject;
        }

        private sealed class GatedUITestResources : FlowScope.Resources.IResourceService
        {
            private readonly string _key;
            private readonly GameObject _asset;
            private readonly TaskCompletionSource<FlowScope.Resources.IResourceHandle<GameObject>> _gate =
                new(TaskCreationOptions.RunContinuationsAsynchronously);

            public GatedUITestResources(string key, GameObject asset)
            {
                _key = key;
                _asset = asset;
                Handle = new UITestResourceHandle<GameObject>(key, asset);
            }

            public int LoadCalls { get; private set; }
            public UITestResourceHandle<GameObject> Handle { get; }
            public bool LoadCancellationTokenCanBeCanceled { get; private set; }

            public Task<FlowScope.Resources.IResourceHandle<T>> LoadAsync<T>(
                string key,
                CancellationToken cancellationToken)
                where T : class
            {
                cancellationToken.ThrowIfCancellationRequested();
                LoadCancellationTokenCanBeCanceled = cancellationToken.CanBeCanceled;
                if (cancellationToken.CanBeCanceled)
                {
                    cancellationToken.Register(() => _gate.TrySetCanceled(cancellationToken));
                }

                LoadCalls++;
                Assert.That(key, Is.EqualTo(_key));
                return AwaitHandle<T>();
            }

            public FlowScope.Resources.IResourceGroup CreateGroup()
            {
                throw new System.NotSupportedException();
            }

            public void Complete()
            {
                _gate.TrySetResult(Handle);
            }

            private async Task<FlowScope.Resources.IResourceHandle<T>> AwaitHandle<T>()
                where T : class
            {
                var handle = await _gate.Task;
                Assert.That(handle.Asset, Is.SameAs(_asset));
                return (FlowScope.Resources.IResourceHandle<T>)(object)handle;
            }
        }

        private sealed class SequencedGatedUITestResources : FlowScope.Resources.IResourceService
        {
            private readonly string _key;
            private readonly List<GameObject> _assets = new();
            private readonly List<UITestResourceHandle<GameObject>> _handles = new();
            private readonly List<TaskCompletionSource<FlowScope.Resources.IResourceHandle<GameObject>>> _gates = new();

            public SequencedGatedUITestResources(string key, params GameObject[] assets)
            {
                _key = key;
                foreach (var asset in assets)
                {
                    _assets.Add(asset);
                    _handles.Add(new UITestResourceHandle<GameObject>(key, asset));
                    _gates.Add(new TaskCompletionSource<FlowScope.Resources.IResourceHandle<GameObject>>(
                        TaskCreationOptions.RunContinuationsAsynchronously));
                }
            }

            public int LoadCalls { get; private set; }

            public UITestResourceHandle<GameObject> HandleAt(int index)
            {
                return _handles[index];
            }

            public Task<FlowScope.Resources.IResourceHandle<T>> LoadAsync<T>(
                string key,
                CancellationToken cancellationToken)
                where T : class
            {
                cancellationToken.ThrowIfCancellationRequested();
                Assert.That(key, Is.EqualTo(_key));

                var loadIndex = LoadCalls;
                LoadCalls++;
                Assert.That(loadIndex, Is.LessThan(_gates.Count), "Unexpected extra resource load.");
                return AwaitHandle<T>(loadIndex);
            }

            public FlowScope.Resources.IResourceGroup CreateGroup()
            {
                throw new System.NotSupportedException();
            }

            public void Complete(int index)
            {
                _gates[index].TrySetResult(_handles[index]);
            }

            private async Task<FlowScope.Resources.IResourceHandle<T>> AwaitHandle<T>(int index)
                where T : class
            {
                var handle = await _gates[index].Task;
                Assert.That(handle.Asset, Is.SameAs(_assets[index]));
                return (FlowScope.Resources.IResourceHandle<T>)(object)handle;
            }
        }
    }
}