From 9269b5009b75d0b5ae3c0e260b0892bc355bbdb7 Mon Sep 17 00:00:00 2001
From: Shane Lee <shane.d.lee@gmail.com>
Date: Thu, 4 Apr 2024 10:06:08 -0600
Subject: [PATCH] Fix tests... maybe

---
 salt/modules/chocolatey.py                    | 53 +++++++++++--------
 salt/states/chocolatey.py                     |  4 +-
 .../states/chocolatey/test_bootstrap.py       |  7 +--
 3 files changed, 36 insertions(+), 28 deletions(-)

diff --git a/salt/modules/chocolatey.py b/salt/modules/chocolatey.py
index d618a175572..64b767c12e2 100644
--- a/salt/modules/chocolatey.py
+++ b/salt/modules/chocolatey.py
@@ -90,10 +90,22 @@ def _no_progress():
     return __context__["chocolatey._no_progress"]
 
 
-def _find_chocolatey():
+def _find_chocolatey(refresh=False):
     """
-    Returns the full path to chocolatey.bat on the host.
+    Returns the full path to the chocolatey binary on the host. If found, the
+    location is cached in ``__context__``.
+
+    Args:
+
+        refresh (bool):
+            Refresh the cached location of the chocolatey binary in
+            ``__context__``
+
+            .. versionadded:: 3007.1
     """
+    if refresh:
+        __context__.pop("chocolatey._path", False)
+
     # Check context
     if "chocolatey._path" in __context__:
         return __context__["chocolatey._path"]
@@ -136,7 +148,7 @@ def chocolatey_version(refresh=False):
         refresh (bool):
             Refresh the cached version of chocolatey
 
-            .. versionadded:: 3008.0
+            .. versionadded:: 3007.1
 
     CLI Example:
 
@@ -150,8 +162,7 @@ def chocolatey_version(refresh=False):
     if "chocolatey._version" in __context__:
         return __context__["chocolatey._version"]
 
-    cmd = [_find_chocolatey()]
-    cmd.append("-v")
+    cmd = [_find_chocolatey(refresh=refresh), "-v"]
     out = __salt__["cmd.run"](cmd, python_shell=False)
     __context__["chocolatey._version"] = out
 
@@ -202,7 +213,7 @@ def bootstrap(force=False, source=None, version=None):
             The version of chocolatey to install. The latest version is
             installed if this value is ``None``. Default is ``None``
 
-            .. versionadded:: 3008.0
+            .. versionadded:: 3007.1
 
     Returns:
         str: The stdout of the Chocolatey installation script
@@ -404,24 +415,24 @@ def unbootstrap():
         salt * chocolatey.unbootstrap
     """
     removed = []
-
-    # Delete the Chocolatey directory
-    choco_dir = os.environ.get("ChocolateyInstall", False)
-    if choco_dir:
-        if os.path.exists(choco_dir):
-            log.debug("Removing Chocolatey directory: %s", choco_dir)
-            __salt__["file.remove"](path=choco_dir, force=True)
-            removed.append(f"Removed Directory: {choco_dir}")
-    else:
-        known_paths = [
+    known_paths = []
+    # Get the install location from the registry first
+    if os.environ.get("ChocolateyInstall", False):
+        known_paths.append(os.environ.get("ChocolateyInstall"))
+    known_paths.extend(
+        [
+            # Default location
             os.path.join(os.environ.get("ProgramData"), "Chocolatey"),
+            # Old default location
             os.path.join(os.environ.get("SystemDrive"), "Chocolatey"),
         ]
-        for path in known_paths:
-            if os.path.exists(path):
-                log.debug("Removing Chocolatey directory: %s", path)
-                __salt__["file.remove"](path=path, force=True)
-                removed.append(f"Removed Directory: {path}")
+    )
+    # Delete all known Chocolatey directories
+    for path in known_paths:
+        if os.path.exists(path):
+            log.debug("Removing Chocolatey directory: %s", path)
+            __salt__["file.remove"](path=path, force=True)
+            removed.append(f"Removed Directory: {path}")
 
     # Delete all Chocolatey environment variables
     for env_var in __salt__["environ.items"]():
diff --git a/salt/states/chocolatey.py b/salt/states/chocolatey.py
index 9c8f1951cad..e84bbbb6a5e 100644
--- a/salt/states/chocolatey.py
+++ b/salt/states/chocolatey.py
@@ -517,7 +517,7 @@ def source_present(
 
 def bootstrapped(name, force=False, source=None, version=None):
     """
-    .. versionadded:: 3008.0
+    .. versionadded:: 3007.1
 
     Ensure chocolatey is installed on the system.
 
@@ -641,7 +641,7 @@ def bootstrapped(name, force=False, source=None, version=None):
 
 def unbootstrapped(name):
     """
-    .. versionadded:: 3008.0
+    .. versionadded:: 3007.1
 
     Ensure chocolatey is removed from the system.
 
diff --git a/tests/pytests/functional/states/chocolatey/test_bootstrap.py b/tests/pytests/functional/states/chocolatey/test_bootstrap.py
index 1a947dd1ff0..42879a843e5 100644
--- a/tests/pytests/functional/states/chocolatey/test_bootstrap.py
+++ b/tests/pytests/functional/states/chocolatey/test_bootstrap.py
@@ -1,5 +1,3 @@
-import logging
-
 import pytest
 
 from salt.exceptions import CommandExecutionError
@@ -11,8 +9,6 @@ pytestmark = [
     pytest.mark.windows_whitelisted,
 ]
 
-log = logging.getLogger(__name__)
-
 
 @pytest.fixture(scope="module")
 def chocolatey(states):
@@ -56,8 +52,9 @@ def clean(chocolatey_mod):
 
 @pytest.fixture()
 def installed(chocolatey_mod):
+    # Even though this variable is not used, it will be displayed in test output
+    # if there is an error
     result = chocolatey_mod.bootstrap(force=True)
-    log.debug("CHOCO: chocolatey install output\n%s", result)
     try:
         chocolatey_version = chocolatey_mod.chocolatey_version(refresh=True)
     except CommandExecutionError: