guard container extension from thrown errors

dansiegel · Nov 24, 2020 · 85c781b · 85c781b
1 parent 113c1b0
commit 85c781b
Showing 1 changed file with 25 additions and 6 deletions.
diff --git a/src/Prism.Container.Extensions/Internals/ContainerLocationHelper.cs b/src/Prism.Container.Extensions/Internals/ContainerLocationHelper.cs
@@ -1,4 +1,5 @@
-using System.ComponentModel;
+using System;
+using System.ComponentModel;
 using System.Linq;
 using System.Reflection;
 using Prism.Ioc;
@@ -16,13 +17,31 @@ public static IContainerExtension LocateContainer(IContainerExtension container
                 ContainerLocator.SetContainerExtension(() => container);
             }
 
-            var located = ContainerLocator.Current;
-            if (located != null)
-                return located;
+            try
+            {
+                var located = ContainerLocator.Current;
+                if (located != null)
+                    return located;
+
+                // If somehow we have an actual null container let's be sure to refresh the Locator
+                ContainerLocator.ResetContainer();
+            }
+            catch
+            {
+                // suppress any errors
+                ContainerLocator.ResetContainer();
+            }
 
             var entryAssembly = Assembly.GetEntryAssembly();
-            var containerExtensionAttribute = entryAssembly.GetCustomAttributes().OfType<ContainerExtensionAttribute>().FirstOrDefault();
-            return containerExtensionAttribute?.GetContainer();
+            var containerAttributes = entryAssembly.GetCustomAttributes().OfType<ContainerExtensionAttribute>();
+
+            if (!containerAttributes.Any())
+                throw new InvalidOperationException("An instance of the IContainerExtension has not been registered with the ContainerLocator, and no ContainerExtensionAttribute has been found in the entry assembly.");
+            else if (containerAttributes.Count() > 1)
+                throw new InvalidOperationException("More than one ContainerExtensionAttribute has been found on the entry assembly. Only a single ContainerExtension should be referenced.");
+
+            var containerExtensionAttribute = containerAttributes.First();
+            return containerExtensionAttribute.GetContainer();
         }
     }
 }