Clean PR

Author: Manikanta Nallagatla

src/WebJobs.Script.WebHost/Controllers/InstanceController.cs

@@ -50,9 +50,15 @@ public async Task<IActionResult> Assign([FromBody] EncryptedHostAssignmentContex
         [Authorize(Policy = PolicyNames.AdminAuthLevel)]
         public async Task<IActionResult> Assign2([FromBody] FunctionsWorkerContainerAssignmentContext workerAssignmentContext)
         {
+            if (workerAssignmentContext?.AssignmentContext == null || workerAssignmentContext.AssignmentContext == null)
+            {
+                return BadRequest("Assignment context is missing.");
+            }
+
             _logger.LogDebug($"Starting container assignment for host : {Request?.Host}");
+            var assignmentContext = _startupContextProvider.SetContext(workerAssignmentContext.AssignmentContext);
 
-            return await AssignInternal(workerAssignmentContext.AssignmentContext);
+            return await AssignInternal(assignmentContext);
         }
 
         [HttpGet]
@@ -73,10 +79,6 @@ public IActionResult GetHttpHealthStatus()
 
         private async Task<IActionResult> AssignInternal(HostAssignmentContext assignmentContext)
         {
-            if (assignmentContext == null)
-            {
-                _logger.LogDebug($"Assignment context is null");
-            }
             // before starting the assignment we want to perform as much
             // up front validation on the context as possible
             string error = await _instanceManager.ValidateContext(assignmentContext);

src/WebJobs.Script.WebHost/Models/FunctionsWorkerContainerAssignmentContext.cs

@@ -9,6 +9,6 @@ namespace Microsoft.Azure.WebJobs.Script.WebHost.Models
     public class FunctionsWorkerContainerAssignmentContext
     {
         [JsonProperty("assignmentContext")]
-        public HostAssignmentContext AssignmentContext { get; private set; }
+        public HostAssignmentContext? AssignmentContext { get; private set; }
     }
 }

src/WebJobs.Script.WebHost/StartupContextProvider.cs

@@ -142,6 +142,11 @@ public virtual HostAssignmentContext SetContext(EncryptedHostAssignmentContext e
             string decryptedContext = EncryptionHelper.Decrypt(encryptedContext.EncryptedContext, environment: _environment);
             var hostAssignmentContext = JsonConvert.DeserializeObject<HostAssignmentContext>(decryptedContext);
 
+            return SetContext(hostAssignmentContext);
+        }
+
+        public virtual HostAssignmentContext SetContext(HostAssignmentContext hostAssignmentContext)
+        {
             // Don't update StartupContext for warmup requests
             if (!hostAssignmentContext.IsWarmupRequest)
             {