Improve first run -y behaviour and add tests.

Author: zooba

src/manage/firstrun.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import time
+import winreg
 
 from . import logging
 from .pathutils import Path
@@ -51,14 +52,20 @@ def check_app_alias(cmd):
     LOGGER.debug("Check passed: aliases are correct")
     return True
 
+_LONG_PATH_KEY = r"System\CurrentControlSet\Control\FileSystem"
+_LONG_PATH_VALUENAME = "LongPathsEnabled"
 
-def check_long_paths(cmd):
+def check_long_paths(
+    cmd,
+    *,
+    hive=winreg.HKEY_LOCAL_MACHINE,
+    keyname=_LONG_PATH_KEY,
+    valuename=_LONG_PATH_VALUENAME,
+):
     LOGGER.debug("Checking long paths setting")
-    import winreg
     try:
-        with winreg.OpenKeyEx(winreg.HKEY_LOCAL_MACHINE,
-                              r"System\CurrentControlSet\Control\FileSystem") as key:
-            if winreg.QueryValueEx(key, "LongPathsEnabled") == (1, winreg.REG_DWORD):
+        with winreg.OpenKeyEx(hive, keyname) as key:
+            if winreg.QueryValueEx(key, valuename) == (1, winreg.REG_DWORD):
                 LOGGER.debug("Check passed: registry key is OK")
                 return True
     except FileNotFoundError:
@@ -67,6 +74,41 @@ def check_long_paths(cmd):
     return False
 
 
+def do_configure_long_paths(
+    cmd,
+    *,
+    hive=winreg.HKEY_LOCAL_MACHINE,
+    keyname=_LONG_PATH_KEY,
+    valuename=_LONG_PATH_VALUENAME,
+):
+    LOGGER.debug("Updating long paths setting")
+    try:
+        with winreg.CreateKeyEx(hive, keyname) as key:
+            winreg.SetValueEx(key, valuename, None, winreg.REG_DWORD, 1)
+            LOGGER.info("The setting has been successfully updated, and will "
+                        "take effect after the next reboot.")
+            return
+    except OSError:
+        pass
+    if not cmd.confirm:
+        # Without confirmation, we assume we can't elevate, so attempt
+        # as the current user and if it fails just print a message.
+        LOGGER.warn("The setting has not been updated. Please rerun '!B!py "
+                    "install --configure!W! with administrative privileges.")
+        return
+    os.startfile(sys.executable, "runas", "**configure-long-paths", show_cmd=0)
+    for _ in range(5):
+        time.sleep(0.25)
+        if check_long_paths(cmd):
+            LOGGER.info("The setting has been successfully updated, and will "
+                        "take effect after the next reboot.")
+            break
+    else:
+        LOGGER.warn("The setting may not have been updated. Please "
+                    "visit the additional help link at the end for "
+                    "more assistance.")
+
+
 def check_py_on_path(cmd):
     LOGGER.debug("Checking for legacy py.exe on PATH")
     from _native import read_alias_package
@@ -120,7 +162,6 @@ def check_global_dir(cmd):
 
 
 def _check_global_dir_registry(cmd):
-    import winreg
     with winreg.OpenKeyEx(winreg.HKEY_CURRENT_USER, "Environment") as key:
         path, kind = winreg.QueryValueEx(key, "Path")
     LOGGER.debug("Current registry path: %s", path)
@@ -139,7 +180,6 @@ def _check_global_dir_registry(cmd):
 
 
 def do_global_dir_on_path(cmd):
-    import winreg
     added = notified = False
     try:
         LOGGER.debug("Adding %s to PATH", cmd.global_dir)
@@ -290,17 +330,8 @@ def first_run(cmd):
                          "may need an administrator to approve, and will require a "
                          "reboot. Some packages may fail to install without long "
                          "path support enabled.\n", wrap=True)
-            if cmd.confirm and not cmd.ask_ny("Update setting now?"):
-                os.startfile(sys.executable, "runas", "**configure-long-paths", show_cmd=0)
-                for _ in range(5):
-                    time.sleep(0.25)
-                    if check_long_paths(cmd):
-                        LOGGER.info("The setting has been successfully updated.")
-                        break
-                else:
-                    LOGGER.warn("The setting may not have been updated. Please "
-                                "visit the additional help link at the end for "
-                                "more assistance.")
+            if not cmd.confirm or not cmd.ask_ny("Update setting now?"):
+                do_configure_long_paths(cmd)
         elif cmd.explicit:
             LOGGER.info("Checked system long paths setting")
 
@@ -314,10 +345,7 @@ def first_run(cmd):
             LOGGER.print("\nThis may interfere with launching the new 'py' "
                          "command, and may be resolved by uninstalling "
                          "'!B!Python launcher!W!'.\n", wrap=True)
-            if (
-                cmd.confirm and
-                not cmd.ask_ny("Open Installed apps now?")
-            ):
+            if cmd.confirm and not cmd.ask_ny("Open Installed apps now?"):
                 os.startfile("ms-settings:appsfeatures")
         elif cmd.explicit:
             if r == "skip":
@@ -342,7 +370,7 @@ def first_run(cmd):
                          "must manually edit environment variables to later "
                          "remove the entry.\n", wrap=True)
             if (
-                cmd.confirm and
+                not cmd.confirm or
                 not cmd.ask_ny("Add commands directory to your PATH now?")
             ):
                 do_global_dir_on_path(cmd)
@@ -352,7 +380,7 @@ def first_run(cmd):
             else:
                 LOGGER.info("Checked PATH for versioned commands directory")
 
-    # This check must be last, because 'do_install' will exit the program.
+    # This check must be last, because a failed install may exit the program.
     if cmd.check_any_install:
         if not check_any_install(cmd):
             welcome()
@@ -365,7 +393,7 @@ def first_run(cmd):
                          "install, or one will be installed automatically when "
                          "needed.\n", wrap=True)
             LOGGER.info("")
-            if cmd.ask_yn("Install CPython now?"):
+            if not cmd.confirm or cmd.ask_yn("Install CPython now?"):
                 do_install(cmd)
         elif cmd.explicit:
             LOGGER.info("Checked for any Python installs")

tests/conftest.py

@@ -148,6 +148,7 @@ def localserver():
 class FakeConfig:
     def __init__(self, global_dir, installs=[]):
         self.global_dir = global_dir
+        self.confirm = False
         self.installs = list(installs)
         self.shebang_can_run_anything = True
         self.shebang_can_run_anything_silently = False

tests/test_firstrun.py

@@ -237,3 +237,27 @@ def test_do_global_dir_path_fail_broadcast(protect_reg, fake_config, assert_log,
     monkeypatch.setattr(winreg, "SetValueEx", lambda *a: None)
     firstrun.do_global_dir_on_path(fake_config)
     assert_log(assert_log.skip_until("Failed to notify of PATH environment.+"))
+
+
+def test_check_long_paths(registry, fake_config):
+    assert not firstrun.check_long_paths(fake_config, hive=registry.hive, keyname=registry.root)
+    registry.setup(LongPathsEnabled=1)
+    assert firstrun.check_long_paths(fake_config, hive=registry.hive, keyname=registry.root)
+
+
+def test_do_configure_long_paths(registry, fake_config, monkeypatch):
+    monkeypatch.setattr(os, "startfile", _raise_oserror)
+    firstrun.do_configure_long_paths(fake_config, hive=registry.hive, keyname=registry.root)
+    assert winreg.QueryValueEx(registry.key, "LongPathsEnabled") == (1, winreg.REG_DWORD)
+
+
+def test_do_configure_long_paths_elevated(protect_reg, fake_config, monkeypatch):
+    startfile_calls = []
+    def startfile(*a, **kw):
+        startfile_calls.append((a, kw))
+    monkeypatch.setattr(os, "startfile", startfile)
+    # Pretend we can interact, so that os.startfile gets called
+    fake_config.confirm = True
+    firstrun.do_configure_long_paths(fake_config)
+    assert startfile_calls
+    assert startfile_calls[0][0][1:] == ("runas", "**configure-long-paths")