From d3e25db330298af606aec2c9a5818a5acad53aac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Per=20Malm=C3=A9n?= <per.malmen@nordlight.io>
Date: Tue, 23 Apr 2024 11:48:48 +0200
Subject: [PATCH] Added test and fix for when a game object throws during
 instantiation

---
 .../Fixtures/CrashingSampleMonoBehaviour.cs   | 29 +++++++++++++++++++
 .../CrashingSampleMonoBehaviour.cs.meta       |  3 ++
 .../Tests/Unity/UnityObjectResolverTest.cs    | 17 +++++++++++
 .../Unity/ObjectResolverUnityExtensions.cs    | 20 ++++++++-----
 4 files changed, 61 insertions(+), 8 deletions(-)
 create mode 100644 VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs
 create mode 100644 VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs.meta

diff --git a/VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs b/VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs
new file mode 100644
index 00000000..fc44c0b3
--- /dev/null
+++ b/VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs
@@ -0,0 +1,29 @@
+using System;
+using UnityEngine;
+
+namespace VContainer.Tests.Unity
+{
+
+    public sealed class CrashingSampleMonoBehaviour : MonoBehaviour, IComponent
+    {
+        public ServiceA ServiceA;
+        public ServiceA ServiceAInAwake;
+        public bool StartCalled;
+        public int UpdateCalls;
+
+        void Start() => StartCalled = true;
+        void Update() => UpdateCalls += 1;
+
+        void Awake()
+        {
+            ServiceAInAwake = ServiceA;
+        }
+
+        [Inject]
+        public void Construct(ServiceA serviceA)
+        {
+            ServiceA = serviceA;
+            throw new Exception("Something broke during construct");
+        }
+    }
+}
diff --git a/VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs.meta b/VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs.meta
new file mode 100644
index 00000000..c57693b3
--- /dev/null
+++ b/VContainer/Assets/Tests/Unity/Fixtures/CrashingSampleMonoBehaviour.cs.meta
@@ -0,0 +1,3 @@
+fileFormatVersion: 2
+guid: 02481ae92e864544bd3bb0f9df8d099d
+timeCreated: 1713431065
\ No newline at end of file
diff --git a/VContainer/Assets/Tests/Unity/UnityObjectResolverTest.cs b/VContainer/Assets/Tests/Unity/UnityObjectResolverTest.cs
index d1210983..29d44bc6 100644
--- a/VContainer/Assets/Tests/Unity/UnityObjectResolverTest.cs
+++ b/VContainer/Assets/Tests/Unity/UnityObjectResolverTest.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Reflection;
 using NUnit.Framework;
 using UnityEngine;
 using VContainer.Unity;
@@ -141,6 +142,22 @@ void AssertInstantiatedInstance(SampleMonoBehaviour instance)
                 Assert.That(instance.ServiceA, Is.InstanceOf<ServiceA>());
             }
         }
+        
+        [Test]
+        public void WhenFailingDuringInstantiateGameObject_TheGameObjectIsStillEnabled()
+        {
+            var builder = new ContainerBuilder();
+            builder.Register<ServiceA>(Lifetime.Singleton);
+            var container = builder.Build();
+
+            var parent = new GameObject("Parent");
+            GameObject original = new GameObject("Original");
+            original.AddComponent<CrashingSampleMonoBehaviour>();
+
+            Assert.Catch(() => container.Instantiate(original, parent.transform));
+            
+            Assert.That(original.gameObject.activeSelf, Is.True);
+        }
 
         [Test]
         public void Instantiate_ExceptionThrownDuringInjection_PrefabIsStillEnabled()
diff --git a/VContainer/Assets/VContainer/Runtime/Unity/ObjectResolverUnityExtensions.cs b/VContainer/Assets/VContainer/Runtime/Unity/ObjectResolverUnityExtensions.cs
index 2d00cb45..fcff370d 100644
--- a/VContainer/Assets/VContainer/Runtime/Unity/ObjectResolverUnityExtensions.cs
+++ b/VContainer/Assets/VContainer/Runtime/Unity/ObjectResolverUnityExtensions.cs
@@ -183,14 +183,18 @@ public static GameObject Instantiate(this IObjectResolver resolver, GameObject p
             var wasActive = prefab.activeSelf;
             prefab.SetActive(false);
 
-            var instance = UnityEngine.Object.Instantiate(prefab, parent, worldPositionStays);
-            SetName(instance, prefab);
-
-            resolver.InjectGameObject(instance);
-
-            prefab.SetActive(wasActive);
-            instance.SetActive(wasActive);
-
+            GameObject instance = null;
+            try
+            {
+                instance = UnityEngine.Object.Instantiate(prefab, parent, worldPositionStays);
+                SetName(instance, prefab);
+                resolver.InjectGameObject(instance);
+            }
+            finally
+            {
+                prefab.SetActive(wasActive);
+                instance?.SetActive(wasActive);    
+            }
             return instance;
         }