Refactor site creation logic to improve timeout handling and ensure proper site properties update

Author: gautamdsheth

src/Commands/Admin/NewTenantSite.cs

@@ -1,9 +1,11 @@
-using Microsoft.Online.SharePoint.TenantManagement;
+using Microsoft.Online.SharePoint.TenantAdministration;
+using Microsoft.Online.SharePoint.TenantManagement;
 using Microsoft.SharePoint.Client;
 using PnP.Framework;
 using PnP.PowerShell.Commands.Base;
 using System;
 using System.Management.Automation;
+using System.Threading.Tasks;
 
 namespace PnP.PowerShell.Commands
 {
@@ -59,23 +61,80 @@ protected override void ExecuteCmdlet()
                 Uri uri = BaseUri;
                 Url = $"{uri.ToString().TrimEnd('/')}/{Url.TrimStart('/')}";
             }
-            Func<TenantOperationMessage, bool> timeoutFunction = TimeoutFunction;
 
-            Guid newSiteId = Tenant.CreateSiteCollection(Url, Title, Owner, Template, (int)StorageQuota,
-                (int)StorageQuotaWarningLevel, TimeZone, (int)ResourceQuota, (int)ResourceQuotaWarningLevel, Lcid,
-                RemoveDeletedSite, Wait, Wait == true ? timeoutFunction : null);
-
-            if (newSiteId != Guid.Empty && Wait && SharingCapability.HasValue)
+            if (ParameterSpecified(nameof(RemoveDeletedSite)))
             {
-                var props = Tenant.GetSitePropertiesByUrl(Url, true);
-                Tenant.Context.Load(props);
-                Tenant.Context.ExecuteQueryRetry();
+                Func<TenantOperationMessage, bool> timeoutFunction = TimeoutFunction;
+
+                Guid newSiteId = Tenant.CreateSiteCollection(Url, Title, Owner, Template, (int)StorageQuota,
+                    (int)StorageQuotaWarningLevel, TimeZone, (int)ResourceQuota, (int)ResourceQuotaWarningLevel, Lcid,
+                    RemoveDeletedSite, Wait, Wait == true ? timeoutFunction : null);
+
+                if (newSiteId != Guid.Empty && Wait && SharingCapability.HasValue)
+                {
+                    var props = Tenant.GetSitePropertiesByUrl(Url, true);
+                    Tenant.Context.Load(props);
+                    Tenant.Context.ExecuteQueryRetry();
 
-                props.SharingCapability = SharingCapability.Value;
+                    props.SharingCapability = SharingCapability.Value;
 
-                var op = props.Update();
-                AdminContext.Load(op, i => i.IsComplete, i => i.PollingInterval);
+                    var op = props.Update();
+                    AdminContext.Load(op, i => i.IsComplete, i => i.PollingInterval);
+                    AdminContext.ExecuteQueryRetry();
+                }
+            }
+            else
+            {
+                SiteCreationProperties siteCreationProperties = new SiteCreationProperties
+                {
+                    Url = Url,
+                    Owner = Owner,
+                    Title = Title,
+                    Template = Template,
+                    StorageMaximumLevel = StorageQuota,
+                    StorageWarningLevel = StorageQuotaWarningLevel,
+                    TimeZoneId = TimeZone,
+                    UserCodeMaximumLevel = ResourceQuota,
+                    UserCodeWarningLevel = ResourceQuotaWarningLevel,
+                    Lcid = Lcid
+                };
+
+                SpoOperation spoOperation = Tenant.CreateSite(siteCreationProperties);
+                AdminContext.Load(spoOperation, s => s.IsComplete, s => s.PollingInterval, s => s.HasTimedout);
                 AdminContext.ExecuteQueryRetry();
+
+                if (Wait)
+                {
+                    while (!spoOperation.IsComplete)
+                    {
+                        if (spoOperation.HasTimedout)
+                        {
+                            throw new TimeoutException("Wait for site creation operation to complete has timed out.");
+                        }
+                        Task.Delay(TimeSpan.FromMilliseconds(spoOperation.PollingInterval)).GetAwaiter().GetResult();                        
+                        spoOperation.RefreshLoad();
+                        if (((Cmdlet)this).Stopping)
+                        {
+                            ((Cmdlet)this).WriteWarning("Cmdlet execution interrupted by user, stopping wait for site creation operation to complete.");
+                            break;
+                        }
+                        AdminContext.Load(spoOperation, s => s.IsComplete, s => s.PollingInterval, s => s.HasTimedout);
+                        AdminContext.ExecuteQueryRetry();
+                    }
+                }
+
+                if (SharingCapability.HasValue)
+                {
+                    var props = Tenant.GetSitePropertiesByUrl(Url, true);
+                    Tenant.Context.Load(props);
+                    Tenant.Context.ExecuteQueryRetry();
+
+                    props.SharingCapability = SharingCapability.Value;
+
+                    var op = props.Update();
+                    AdminContext.Load(op, i => i.IsComplete, i => i.PollingInterval);
+                    AdminContext.ExecuteQueryRetry();
+                }
             }
         }