From fe9bdaadd2fe39aeae60a9daeb8613b823dddce1 Mon Sep 17 00:00:00 2001
From: Joni Savolainen <joni@savolainen.io>
Date: Sun, 9 Oct 2022 10:28:24 +0300
Subject: [PATCH] feat(UIComponent): wait for child components to be
 initialized

BREAKING CHANGE: UIComponents now wait for their hierarchy to be initialized before calling OnInit.
---
 Assets/UIComponents.Tests/UIComponentTests.cs | 66 +++++++++++++++++--
 Assets/UIComponents/Core/UIComponent.cs       | 12 ++++
 2 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/Assets/UIComponents.Tests/UIComponentTests.cs b/Assets/UIComponents.Tests/UIComponentTests.cs
index b1df6958..27942512 100644
--- a/Assets/UIComponents.Tests/UIComponentTests.cs
+++ b/Assets/UIComponents.Tests/UIComponentTests.cs
@@ -18,19 +18,23 @@ public class Initialization
         {
             private class MockAssetResolver : IAssetResolver
             {
-                public readonly Dictionary<string, object> TaskCompletionSources =
+                private readonly Dictionary<string, object> _taskCompletionSources =
                     new Dictionary<string, object>();
 
                 public Task<T> LoadAsset<T>(string assetPath) where T : Object
                 {
-                    var source = new TaskCompletionSource<T>();
-                    TaskCompletionSources.Add(assetPath, source);
-                    return source.Task;
+                    if (!_taskCompletionSources.ContainsKey(assetPath))
+                    {
+                        var source = new TaskCompletionSource<T>();
+                        _taskCompletionSources.Add(assetPath, source);
+                    }
+
+                    return ((TaskCompletionSource<T>) _taskCompletionSources[assetPath]).Task;
                 }
 
                 public void CompleteLoad<T>(string assetPath) where T : ScriptableObject
                 {
-                    var source = (TaskCompletionSource<T>) TaskCompletionSources[assetPath];
+                    var source = (TaskCompletionSource<T>) _taskCompletionSources[assetPath];
                     source.SetResult(ScriptableObject.CreateInstance<T>());
                 }
 
@@ -117,6 +121,58 @@ public IEnumerator Allows_Waiting_For_Initialization_With_Obsolete_Method()
                 Assert.That(component.Initialized, Is.True);
             }
             
+            [Layout("Child")]
+            [Stylesheet("ChildStylesheet")]
+            private class ChildComponent : UIComponent {}
+            
+            [Layout("NestedChild")]
+            private class NestedChildComponent : UIComponent {}
+
+            [Test]
+            public void Does_Not_Initialize_If_Children_Are_Uninitialized()
+            {
+                var component = _testBed.CreateComponent<TestComponent>();
+
+                var firstChild = _testBed.CreateComponent<ChildComponent>();
+                var secondChild = _testBed.CreateComponent<ChildComponent>();
+
+                component.Add(firstChild);
+                component.Add(secondChild);
+                
+                _mockAssetResolver.CompleteLoad<VisualTreeAsset>("Layout");
+                _mockAssetResolver.CompleteLoad<StyleSheet>("Stylesheet1");
+                _mockAssetResolver.CompleteLoad<StyleSheet>("Stylesheet2");
+
+                Assert.That(component.Initialized, Is.False);
+            }
+
+            [Test]
+            public void Initializes_When_Children_Are_Initialized()
+            {
+                var component = _testBed.CreateComponent<TestComponent>();
+
+                var firstChild = _testBed.CreateComponent<ChildComponent>();
+                var secondChild = _testBed.CreateComponent<ChildComponent>();
+
+                var nestedChild = _testBed.CreateComponent<NestedChildComponent>();
+                
+                firstChild.Add(nestedChild);
+
+                component.Add(firstChild);
+                component.Add(secondChild);
+                
+                _mockAssetResolver.CompleteLoad<VisualTreeAsset>("Layout");
+                _mockAssetResolver.CompleteLoad<StyleSheet>("Stylesheet1");
+                _mockAssetResolver.CompleteLoad<StyleSheet>("Stylesheet2");
+
+                _mockAssetResolver.CompleteLoad<VisualTreeAsset>("Child");
+                _mockAssetResolver.CompleteLoad<StyleSheet>("ChildStylesheet");
+                
+                _mockAssetResolver.CompleteLoad<VisualTreeAsset>("NestedChild");
+                
+                Assert.That(component.Initialized, Is.True);
+            }
+            
             private class BareTestComponent : UIComponent {}
 
             [Test]
diff --git a/Assets/UIComponents/Core/UIComponent.cs b/Assets/UIComponents/Core/UIComponent.cs
index 788f04e9..f5978bf6 100644
--- a/Assets/UIComponents/Core/UIComponent.cs
+++ b/Assets/UIComponents/Core/UIComponent.cs
@@ -117,6 +117,18 @@ private async void Initialize()
             LoadLayout(layoutAsset);
             LoadStyles(styles);
 
+            var childInitializationTasks = new List<Task>();
+
+            for (var i = 0; i < childCount; i++)
+            {
+                var child = hierarchy.ElementAt(i);
+                
+                if (child is UIComponent component)
+                    childInitializationTasks.Add(component.InitializationTask);
+            }
+
+            await Task.WhenAll(childInitializationTasks);
+
             PostHierarchySetupProfilerMarker.Begin();
 
             ApplyEffects();