Add test for explicitly elevating between detect and plan.

2 files changed, 61 insertions, 0 deletions

diff --git a/src/TestBA/MessagePump.cs b/src/TestBA/MessagePump.cs
new file mode 100644
index 00000000..21a00349
--- /dev/null
+++ b/src/TestBA/MessagePump.cs
@@ -0,0 +1,39 @@
+// Copyright (c) .NET Foundation and contributors. All rights reserved. Licensed under the Microsoft Reciprocal License. See LICENSE.TXT file in the project root for full license information.
+
+namespace WixToolset.Test.BA
+{
+    using System;
+    using System.Runtime.InteropServices;
+    using System.Windows.Forms;
+
+    public class MessagePump
+    {
+        const uint PM_REMOVE = 1;
+
+        [DllImport("user32.dll", ExactSpelling = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        public static extern bool PeekMessageW(ref Message pMsg, IntPtr hWnd, uint wMsgFilterMin, uint wMsgFilterMax, uint wRemoveMsg);
+
+        [DllImport("user32.dll", ExactSpelling = true)]
+        [return: MarshalAs(UnmanagedType.Bool)]
+        public static extern bool TranslateMessage(ref Message pMsg);
+
+        [DllImport("user32.dll", ExactSpelling = true)]
+        public static extern IntPtr DispatchMessageW(ref Message pMsg);
+
+        public static void ProcessMessages(int maxMessages)
+        {
+            for (int i = 0; i < maxMessages; i++)
+            {
+                Message message = new Message();
+                if (!PeekMessageW(ref message, IntPtr.Zero, 0, 0, PM_REMOVE))
+                {
+                    break;
+                }
+
+                TranslateMessage(ref message);
+                DispatchMessageW(ref message);
+            }
+        }
+    }
+}
diff --git a/src/TestBA/TestBA.cs b/src/TestBA/TestBA.cs
index 1348ce98..b9f869a6 100644
--- a/src/TestBA/TestBA.cs
+++ b/src/TestBA/TestBA.cs
@@ -28,6 +28,7 @@ namespace WixToolset.Test.BA
 
         private bool immediatelyQuit;
         private bool quitAfterDetect;
+        private bool explicitlyElevateAndPlanFromOnElevateBegin;
         private int redetectRemaining;
         private int sleepDuringCache;
         private int cancelCacheAtProgress;
@@ -122,6 +123,12 @@ namespace WixToolset.Test.BA
                 redetectCount = 0;
             }
 
+            string explicitlyElevateAndPlanFromOnElevateBegin = this.ReadPackageAction(null, "ExplicitlyElevateAndPlanFromOnElevateBegin");
+            if (String.IsNullOrEmpty(explicitlyElevateAndPlanFromOnElevateBegin) || !Boolean.TryParse(explicitlyElevateAndPlanFromOnElevateBegin, out this.explicitlyElevateAndPlanFromOnElevateBegin))
+            {
+                this.explicitlyElevateAndPlanFromOnElevateBegin = false;
+            }
+
             string quitAfterDetect = this.ReadPackageAction(null, "QuitAfterDetect");
             if (String.IsNullOrEmpty(quitAfterDetect) || !Boolean.TryParse(quitAfterDetect, out this.quitAfterDetect))
             {
@@ -214,6 +221,10 @@ namespace WixToolset.Test.BA
                 {
                     this.ShutdownUiThread();
                 }
+                else if (this.explicitlyElevateAndPlanFromOnElevateBegin)
+                {
+                    this.Engine.Elevate(this.windowHandle);
+                }
                 else
                 {
                     this.Engine.Plan(this.action);
@@ -225,6 +236,17 @@ namespace WixToolset.Test.BA
             }
         }
 
+        protected override void OnElevateBegin(ElevateBeginEventArgs args)
+        {
+            if (this.explicitlyElevateAndPlanFromOnElevateBegin)
+            {
+                this.Engine.Plan(this.action);
+
+                // Simulate showing some UI since these tests won't actually show the UAC prompt.
+                MessagePump.ProcessMessages(10);
+            }
+        }
+
         protected override void OnPlanPackageBegin(PlanPackageBeginEventArgs args)
         {
             RequestState state;