Fixup tests

Author: CyrusNajmabadi

src/EditorFeatures/Core/IntelliSense/AbstractController.cs

@@ -78,11 +78,11 @@ private void OnTextViewClosed(object sender, EventArgs e)
         this.TextView.TextBuffer.PostChanged -= this.OnTextViewBufferPostChanged;
     }
 
-    public Task WaitForModelComputation_ForTestingPurposesOnlyAsync()
+    public Task WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
     {
         this.ThreadingContext.ThrowIfNotOnUIThread();
         VerifySessionIsActive();
-        return sessionOpt.WaitForModelComputation_ForTestingPurposesOnlyAsync();
+        return sessionOpt.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync();
     }
 
     void IController<TModel>.OnModelUpdated(TModel result, bool updateController)

src/EditorFeatures/Core/IntelliSense/ISession.cs

@@ -14,5 +14,5 @@ internal interface ISession<TModel>
 
     void Stop();
 
-    Task WaitForModelComputation_ForTestingPurposesOnlyAsync();
+    Task WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync();
 }

src/EditorFeatures/Core/IntelliSense/ModelComputation.cs

@@ -81,10 +81,19 @@ public Task<TModel> ModelTask
         }
     }
 
-    public Task WaitForModelComputation_ForTestingPurposesOnlyAsync()
-        => ModelTask;
+    public async Task WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
+    {
+        var model = await ModelTask.ConfigureAwait(true);
+        if (!_notifyControllerTask.IsCompleted)
+        {
+            // Reset lastTask so controller.OnModelUpdated is only called once
+            _lastTask = Task.FromResult(model);
+        }
+
+        await _notifyControllerTask.ConfigureAwait(true);
+    }
 
-    public virtual void Stop()
+    public void Stop()
     {
         ThreadingContext.ThrowIfNotOnUIThread();
 

src/EditorFeatures/Core/IntelliSense/Session.cs

@@ -53,9 +53,9 @@ public virtual void Stop()
         this.PresenterSession.Dismiss();
     }
 
-    public Task WaitForModelComputation_ForTestingPurposesOnlyAsync()
+    public Task WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
     {
         Computation.ThreadingContext.ThrowIfNotOnUIThread();
-        return Computation.WaitForModelComputation_ForTestingPurposesOnlyAsync();
+        return Computation.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync();
     }
 }

src/EditorFeatures/Test2/IntelliSense/ModelTests.vb

@@ -56,13 +56,14 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             Dim model = New Model()
             Dim controller = New Mock(Of IController(Of Model))(MockBehavior.Strict)
             controller.Setup(Function(c) c.BeginAsyncOperation("", Nothing, It.IsAny(Of String), It.IsAny(Of Integer))).Returns(EmptyAsyncToken.Instance)
-            controller.Setup(Sub(c) c.OnModelUpdated(model, True))
+            controller.Setup(Sub(c) c.OnModelUpdated(model, False))
+
             Dim modelComputation = TestModelComputation.Create(threadingContext, controller:=controller.Object)
 
             modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c) Task.FromResult(model))
-            Await modelComputation.WaitForModelComputation_ForTestingPurposesOnlyAsync()
+            Await modelComputation.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
 
-            controller.Verify(Sub(c) c.OnModelUpdated(model, True))
+            controller.Verify(Sub(c) c.OnModelUpdated(model, False))
         End Function
 
         <WpfFact>
@@ -83,9 +84,9 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
                                                                       End Function)
             modelComputation.ChainTaskAndNotifyControllerWhenFinished(Function(m, c) Task.FromResult(model))
             Monitor.Exit(gate)
-            Await modelComputation.WaitForModelComputation_ForTestingPurposesOnlyAsync()
+            Await modelComputation.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
 
-            controller.Verify(Sub(c) c.OnModelUpdated(model, True), Times.Once)
+            controller.Verify(Sub(c) c.OnModelUpdated(DirectCast(Nothing, Model), False), Times.Once)
         End Function
 
         <WpfFact>

src/EditorFeatures/Test2/IntelliSense/SignatureHelpControllerTests.vb

@@ -32,7 +32,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
 
             GetMocks(controller).PresenterSession.Setup(Sub(p) p.Dismiss())
 
-            Await controller.WaitForModelComputation_ForTestingPurposesOnlyAsync()
+            Await controller.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
 
             Assert.Equal(0, GetMocks(controller).Provider.GetItemsCount)
         End Function
@@ -80,7 +80,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             Dim controller = Await CreateController(CreateWorkspace(), waitForPresentation:=True)
 
             Mock.Get(GetMocks(controller).View.Object.Caret).Raise(Sub(c) AddHandler c.PositionChanged, Nothing, New CaretPositionChangedEventArgs(Nothing, Nothing, Nothing))
-            Await controller.WaitForModelComputation_ForTestingPurposesOnlyAsync()
+            Await controller.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
 
             ' GetItemsAsync is called once initially, and then once as a result of handling the PositionChanged event
             Assert.Equal(2, GetMocks(controller).Provider.GetItemsCount)
@@ -93,7 +93,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
             DirectCast(controller, IChainedCommandHandler(Of TypeCharCommandArgs)).ExecuteCommand(
                 New TypeCharCommandArgs(CreateMock(Of ITextView), CreateMock(Of ITextBuffer), ")"c),
                 Sub() GetMocks(controller).Buffer.Insert(0, ")"), TestCommandExecutionContext.Create())
-            Await controller.WaitForModelComputation_ForTestingPurposesOnlyAsync()
+            Await controller.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
 
             ' GetItemsAsync is called once initially, and then once as a result of handling the typechar command
             Assert.Equal(2, GetMocks(controller).Provider.GetItemsCount)
@@ -185,7 +185,7 @@ Namespace Microsoft.CodeAnalysis.Editor.UnitTests.IntelliSense
                 DirectCast(controller, IChainedCommandHandler(Of InvokeSignatureHelpCommandArgs)).ExecuteCommand(
                     New InvokeSignatureHelpCommandArgs(view.Object, buffer), Nothing, TestCommandExecutionContext.Create())
                 If waitForPresentation Then
-                    Await controller.WaitForModelComputation_ForTestingPurposesOnlyAsync()
+                    Await controller.WaitForModelComputationAndControllerNotification_ForTestingPurposesOnlyAsync()
                 End If
             End If