From af6252876abb5cf407370201186ee42ab9dbe758 Mon Sep 17 00:00:00 2001 From: Dan Albert Date: Mon, 8 May 2023 23:33:29 -0700 Subject: [PATCH] Factor out wingreg usage to simplify. Querying the registry in each caller made the test mocks quite messy. Factor that out so we can clean up the implementation and the tests. This also removes the "is standalone" and "is steam" APIs since those don't seem useful and they complicated the implementation. The other cleanup this makes easy is switching to the contextmanager OpenKey, which fixes the bug I found when writing the test in the previous PR. --- dcs/installation.py | 171 +++++++--------------- dcs/winreg.py | 25 ++++ tests/test_installation.py | 286 +++++++++++-------------------------- 3 files changed, 154 insertions(+), 328 deletions(-) create mode 100644 dcs/winreg.py diff --git a/dcs/installation.py b/dcs/installation.py index ecde600d..af3a7435 100644 --- a/dcs/installation.py +++ b/dcs/installation.py @@ -6,95 +6,40 @@ (could be done either through windows registry, or through filesystem analysis ?) """ -import sys import os import re +import sys +from pathlib import Path +from typing import List, Optional -is_windows_os = True -try: - import winreg -except ImportError: - is_windows_os = False - print("WARNING : Trying to run pydcs on non Windows machine") - +from dcs.winreg import read_current_user_value -# Note: Steam App ID for DCS World is 223750 -STEAM_REGISTRY_KEY_NAME = "Software\\Valve\\Steam\\Apps\\223750" +STEAM_REGISTRY_KEY_NAME = "Software\\Valve\\Steam" DCS_STABLE_REGISTRY_KEY_NAME = "Software\\Eagle Dynamics\\DCS World" DCS_BETA_REGISTRY_KEY_NAME = "Software\\Eagle Dynamics\\DCS World OpenBeta" -def is_using_dcs_steam_edition(): - """ - Check if DCS World : Steam Edition version is installed on this computer - :return True if DCS Steam edition is installed, - -1 if DCS Steam Edition is registered in Steam apps but not installed, - False if never installed in Steam - """ - if not is_windows_os: - return False - try: - dcs_steam_app_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, STEAM_REGISTRY_KEY_NAME) - installed = winreg.QueryValueEx(dcs_steam_app_key, "Installed") - winreg.CloseKey(dcs_steam_app_key) - if installed[0] == 1: - return True - else: - return False - except FileNotFoundError: - return False - - -def is_using_dcs_standalone_edition(): - """ - Check if DCS World standalone edition is installed on this computer - :return True if Standalone is installed, False if it is not - """ - if not is_windows_os: - return False - try: - dcs_path_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, DCS_STABLE_REGISTRY_KEY_NAME) - winreg.CloseKey(dcs_path_key) - return True - except FileNotFoundError: - try: - dcs_path_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, DCS_BETA_REGISTRY_KEY_NAME) - winreg.CloseKey(dcs_path_key) - return True - except FileNotFoundError: - return False - - -def get_dcs_install_directory(): +def get_dcs_install_directory() -> str: """ Get the DCS World install directory for this computer :return DCS World install directory """ - if is_using_dcs_standalone_edition(): - try: - dcs_path_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Software\\Eagle Dynamics\\DCS World") - path = winreg.QueryValueEx(dcs_path_key, "Path") - dcs_dir = path[0] + os.path.sep - winreg.CloseKey(dcs_path_key) - return dcs_dir - except FileNotFoundError: - try: - dcs_path_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Software\\Eagle Dynamics\\DCS World OpenBeta") - path = winreg.QueryValueEx(dcs_path_key, "Path") - dcs_dir = path[0] + os.path.sep - winreg.CloseKey(dcs_path_key) - return dcs_dir - except FileNotFoundError: - print("Couldn't detect DCS World installation folder", file=sys.stderr) - return "" - except OSError: - print("Couldn't detect DCS World installation folder", file=sys.stderr) - return "" - elif is_using_dcs_steam_edition(): - return _find_steam_dcs_directory() - else: - print("Couldn't detect any installed DCS World version", file=sys.stderr) - return "" + standalone_stable_path = read_current_user_value( + DCS_STABLE_REGISTRY_KEY_NAME, "Path", Path + ) + if standalone_stable_path is not None: + return f"{standalone_stable_path}{os.path.sep}" + standalone_beta_path = read_current_user_value( + DCS_BETA_REGISTRY_KEY_NAME, "Path", Path + ) + if standalone_beta_path is not None: + return f"{standalone_beta_path}{os.path.sep}" + steam_path = _dcs_steam_path() + if steam_path is not None: + return f"{steam_path}{os.path.sep}" + + print("Couldn't detect any installed DCS World version", file=sys.stderr) + return "" def get_dcs_saved_games_directory(): @@ -107,68 +52,50 @@ def get_dcs_saved_games_directory(): if os.path.exists(dcs_variant): # read from the file, append first line to saved games, e.g.: DCS.openbeta with open(dcs_variant, "r") as file: - suffix = re.sub(r'[^\w\d-]', '', file.read()) + suffix = re.sub(r"[^\w\d-]", "", file.read()) saved_games = saved_games + "." + suffix return saved_games -def _find_steam_directory(): - """ - Get the Steam install directory for this computer from registry - :return Steam installation path - """ - try: - steam_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Software\\Valve\\Steam") - path = winreg.QueryValueEx(steam_key, "SteamPath")[0] - winreg.CloseKey(steam_key) - return path - except FileNotFoundError as fnfe: - print(fnfe, file=sys.stderr) - return "" - - -def _get_steam_library_folders(): +def _steam_library_directories() -> List[Path]: """ Get the installation directory for Steam games :return List of Steam library folders where games can be installed """ - try: - steam_dir = _find_steam_directory() - """ - For reference here is what the vdf file is supposed to look like : - - "LibraryFolders" - { - "TimeNextStatsReport" "1561832478" - "ContentStatsID" "-158337411110787451" - "1" "D:\\Games\\Steam" - "2" "E:\\Steam" - } - """ - vdf_file_location = steam_dir + os.path.sep + "steamapps" + os.path.sep + "libraryfolders.vdf" - with open(vdf_file_location) as adf_file: - paths = [line.split("\"")[3] for line in adf_file.readlines()[1:] if ':\\\\' in line] - return paths - except Exception as e: - print(e) + steam_dir = read_current_user_value(STEAM_REGISTRY_KEY_NAME, "SteamPath", Path) + if steam_dir is None: return [] + """ + For reference here is what the vdf file is supposed to look like : + + "LibraryFolders" + { + "TimeNextStatsReport" "1561832478" + "ContentStatsID" "-158337411110787451" + "1" "D:\\Games\\Steam" + "2" "E:\\Steam" + } + """ + contents = (steam_dir / "steamapps/libraryfolders.vdf").read_text() + return [ + Path(line.split('"')[3]) + for line in contents.splitlines()[1:] + if ":\\\\" in line + ] -def _find_steam_dcs_directory(): +def _dcs_steam_path() -> Optional[Path]: """ Find the DCS install directory for DCS World Steam Edition :return: Install directory as string, empty string if not found """ - for library_folder in _get_steam_library_folders(): - folder = library_folder + os.path.sep + "steamapps" + os.path.sep + "common" + os.path.sep + "DCSWorld" - if os.path.isdir(folder): - return folder + os.path.sep - return "" + for library_folder in _steam_library_directories(): + folder = library_folder / "steamapps/common/DCSWorld" + if folder.is_dir(): + return folder + return None if __name__ == "__main__": - print("Using Windows : " + str(is_windows_os)) - print("Using STEAM Edition : " + str(is_using_dcs_steam_edition())) - print("Using Standalone Edition : " + str(is_using_dcs_standalone_edition())) print("DCS Installation directory : " + get_dcs_install_directory()) print("DCS saved games directory : " + get_dcs_saved_games_directory()) diff --git a/dcs/winreg.py b/dcs/winreg.py new file mode 100644 index 00000000..4480b9af --- /dev/null +++ b/dcs/winreg.py @@ -0,0 +1,25 @@ +"""Wrapper around the stdlib winreg that fails gracefully on non-Windows.""" +import logging +import sys +from typing import Any, Callable, Optional, TypeVar + +T = TypeVar("T") + + +def read_current_user_value( + key: str, value: str, ctor: Callable[[Any], T] = lambda x: x +) -> Optional[T]: + if sys.platform == "win32": + import winreg + + try: + with winreg.OpenKey(winreg.HKEY_CURRENT_USER, key) as hkey: + # QueryValueEx returns a tuple of (value, type ID). + return ctor(winreg.QueryValueEx(hkey, value)[0]) + except FileNotFoundError: + return None + else: + logging.getLogger("pydcs").error( + "Cannot read registry keys on non-Windows OS, returning None" + ) + return None diff --git a/tests/test_installation.py b/tests/test_installation.py index 19bd9613..a8ae2bc1 100644 --- a/tests/test_installation.py +++ b/tests/test_installation.py @@ -1,11 +1,8 @@ -import sys - -try: - import winreg -except ImportError: - pass +import textwrap +from collections.abc import Iterator from pathlib import Path -from unittest.mock import Mock, call, patch +from typing import Any, Callable, Optional, TypeVar +from unittest.mock import Mock, patch import pytest @@ -15,227 +12,104 @@ STEAM_REGISTRY_KEY_NAME, get_dcs_install_directory, get_dcs_saved_games_directory, - is_using_dcs_standalone_edition, - is_using_dcs_steam_edition, -) - -pytestmark = pytest.mark.skipif( - sys.platform != "win32", reason="dcs.installation is windows only" ) - -@patch("winreg.CloseKey") -@patch("winreg.QueryValueEx") -@patch("winreg.OpenKey") -def test_is_using_dcs_steam_edition( - mock_openkey: Mock, mock_queryvalueex: Mock, mock_closekey: Mock -) -> None: - mock_openkey.return_value = "key" - mock_queryvalueex.return_value = (1, 4) - - assert is_using_dcs_steam_edition() - mock_openkey.assert_called_once_with( - winreg.HKEY_CURRENT_USER, STEAM_REGISTRY_KEY_NAME - ) - mock_queryvalueex.assert_called_once_with("key", "Installed") - mock_closekey.assert_called_once_with("key") +T = TypeVar("T") -@patch("winreg.CloseKey") -@patch("winreg.OpenKey") -def test_is_using_dcs_steam_edition_no_key( - mock_openkey: Mock, mock_closekey: Mock +def configure_registry_mock( + mock: Mock, + steam: Optional[Path] = None, + standalone_stable: Optional[Path] = None, + standalone_beta: Optional[Path] = None, ) -> None: - mock_openkey.side_effect = FileNotFoundError - - assert not is_using_dcs_steam_edition() - mock_openkey.assert_called_once_with( - winreg.HKEY_CURRENT_USER, STEAM_REGISTRY_KEY_NAME - ) - mock_closekey.assert_not_called() - - -@pytest.mark.xfail(strict=True) # Does not close the key. -@patch("winreg.CloseKey") -@patch("winreg.QueryValueEx") -@patch("winreg.OpenKey") -def test_is_using_dcs_steam_edition_no_value( - mock_openkey: Mock, mock_queryvalueex: Mock, mock_closekey: Mock -) -> None: - mock_openkey.return_value = "key" - mock_queryvalueex.side_effect = FileNotFoundError - - assert not is_using_dcs_steam_edition() - mock_openkey.assert_called_once_with( - winreg.HKEY_CURRENT_USER, STEAM_REGISTRY_KEY_NAME + def registry_mock( + key: str, value: str, ctor: Callable[[Any], T] = lambda x: x + ) -> Optional[T]: + if ( + key == STEAM_REGISTRY_KEY_NAME + and value == "SteamPath" + and steam is not None + ): + return ctor(steam) + if ( + key == DCS_STABLE_REGISTRY_KEY_NAME + and value == "Path" + and standalone_stable is not None + ): + return ctor(standalone_stable) + if ( + key == DCS_BETA_REGISTRY_KEY_NAME + and value == "Path" + and standalone_beta is not None + ): + return ctor(standalone_beta) + return None + + mock.side_effect = registry_mock + + +@pytest.fixture(name="steam_dcs_install") +def steam_dcs_install_fixture(tmp_path: Path) -> Iterator[Path]: + escaped_path = str(tmp_path).replace("\\", "\\\\") + vdf_path = tmp_path / "steamapps/libraryfolders.vdf" + vdf_path.parent.mkdir(parents=True) + vdf_path.write_text( + textwrap.dedent( + f"""\ + "LibraryFolders" + {{ + "TimeNextStatsReport" "1561832478" + "ContentStatsID" "-158337411110787451" + "1" "D:\\\\Games\\\\Steam" + "2" "{escaped_path}" + }} + """ + ) ) - mock_queryvalueex.assert_called_once_with("key", "Installed") - mock_closekey.assert_called_once_with("key") + dcs_install_path = tmp_path / "steamapps/common/DCSWorld" + dcs_install_path.mkdir(parents=True) -@patch("winreg.CloseKey") -@patch("winreg.QueryValueEx") -@patch("winreg.OpenKey") -def test_is_using_dcs_steam_edition_not_installed( - mock_openkey: Mock, mock_queryvalueex: Mock, mock_closekey: Mock -) -> None: - mock_openkey.return_value = "key" - mock_queryvalueex.return_value = (0, 0) + with patch("dcs.installation.read_current_user_value") as mock: + configure_registry_mock(mock, steam=tmp_path) + yield dcs_install_path - assert not is_using_dcs_steam_edition() - mock_openkey.assert_called_once_with( - winreg.HKEY_CURRENT_USER, STEAM_REGISTRY_KEY_NAME - ) - mock_queryvalueex.assert_called_once_with("key", "Installed") - mock_closekey.assert_called_once_with("key") +@pytest.fixture(name="stable_dcs_install") +def stable_dcs_install_fixture(tmp_path: Path) -> Iterator[Path]: + with patch("dcs.installation.read_current_user_value") as mock: + configure_registry_mock(mock, standalone_stable=tmp_path) + yield tmp_path -@patch("winreg.CloseKey") -@patch("winreg.OpenKey") -def test_is_using_dcs_standalone_edition_stable( - mock_openkey: Mock, mock_closekey: Mock -) -> None: - def stable_installed(key_type: int, name: str) -> str: - if ( - key_type == winreg.HKEY_CURRENT_USER - and name == DCS_STABLE_REGISTRY_KEY_NAME - ): - return "key" - raise FileNotFoundError - mock_openkey.side_effect = stable_installed +@pytest.fixture(name="beta_dcs_install") +def beta_dcs_install_fixture(tmp_path: Path) -> Iterator[Path]: + with patch("dcs.installation.read_current_user_value") as mock: + configure_registry_mock(mock, standalone_beta=tmp_path) + yield tmp_path - assert is_using_dcs_standalone_edition() - mock_openkey.assert_has_calls( - [call(winreg.HKEY_CURRENT_USER, DCS_STABLE_REGISTRY_KEY_NAME)] - ) - mock_closekey.assert_called_once_with("key") +@pytest.fixture(name="none_installed") +def none_installed_fixture() -> Iterator[None]: + with patch("dcs.installation.read_current_user_value") as mock: + configure_registry_mock(mock) + yield -@patch("winreg.CloseKey") -@patch("winreg.OpenKey") -def test_is_using_dcs_standalone_edition_beta( - mock_openkey: Mock, mock_closekey: Mock -) -> None: - def stable_installed(key_type: int, name: str) -> str: - if key_type == winreg.HKEY_CURRENT_USER and name == DCS_BETA_REGISTRY_KEY_NAME: - return "key" - raise FileNotFoundError - mock_openkey.side_effect = stable_installed +def test_get_dcs_install_directory_stable(stable_dcs_install: Path) -> None: + assert get_dcs_install_directory() == f"{stable_dcs_install}\\" - assert is_using_dcs_standalone_edition() - mock_openkey.assert_has_calls( - [call(winreg.HKEY_CURRENT_USER, DCS_BETA_REGISTRY_KEY_NAME)] - ) - mock_closekey.assert_called_once_with("key") +def test_get_dcs_install_directory_beta(beta_dcs_install: Path) -> None: + assert get_dcs_install_directory() == f"{beta_dcs_install}\\" -@patch("winreg.CloseKey") -@patch("winreg.OpenKey") -def test_is_using_dcs_standalone_not_installed( - mock_openkey: Mock, mock_closekey: Mock -) -> None: - mock_openkey.side_effect = FileNotFoundError - - assert not is_using_dcs_standalone_edition() - mock_openkey.assert_has_calls( - [ - call(winreg.HKEY_CURRENT_USER, DCS_STABLE_REGISTRY_KEY_NAME), - call(winreg.HKEY_CURRENT_USER, DCS_BETA_REGISTRY_KEY_NAME), - ], - any_order=True, - ) - mock_closekey.assert_not_called() - - -@patch("dcs.installation.is_using_dcs_steam_edition") -@patch("dcs.installation.is_using_dcs_standalone_edition") -@patch("winreg.CloseKey") -@patch("winreg.QueryValueEx") -@patch("winreg.OpenKey") -def test_get_dcs_install_directory_stable( - mock_openkey: Mock, - mock_queryvalueex: Mock, - mock_closekey: Mock, - mock_standalone_edition: Mock, - mock_steam_edition: Mock, -) -> None: - def stable_installed(key_type: int, name: str) -> str: - if ( - key_type == winreg.HKEY_CURRENT_USER - and name == DCS_STABLE_REGISTRY_KEY_NAME - ): - return "key" - raise FileNotFoundError - - mock_openkey.side_effect = stable_installed - mock_queryvalueex.return_value = ("path",) - mock_standalone_edition.return_value = True - mock_steam_edition.return_value = False - - assert get_dcs_install_directory() == "path\\" - - mock_queryvalueex.assert_called_once_with("key", "Path") - mock_closekey.assert_called_once_with("key") - - -@patch("dcs.installation.is_using_dcs_steam_edition") -@patch("dcs.installation.is_using_dcs_standalone_edition") -@patch("winreg.CloseKey") -@patch("winreg.QueryValueEx") -@patch("winreg.OpenKey") -def test_get_dcs_install_directory_beta( - mock_openkey: Mock, - mock_queryvalueex: Mock, - mock_closekey: Mock, - mock_standalone_edition: Mock, - mock_steam_edition: Mock, -) -> None: - def stable_installed(key_type: int, name: str) -> str: - if key_type == winreg.HKEY_CURRENT_USER and name == DCS_BETA_REGISTRY_KEY_NAME: - return "key" - raise FileNotFoundError - - mock_openkey.side_effect = stable_installed - mock_queryvalueex.return_value = ("path",) - mock_standalone_edition.return_value = True - mock_steam_edition.return_value = False - - assert get_dcs_install_directory() == "path\\" - - mock_queryvalueex.assert_called_once_with("key", "Path") - mock_closekey.assert_called_once_with("key") - - -@patch("dcs.installation._get_steam_library_folders") -@patch("dcs.installation.is_using_dcs_steam_edition") -@patch("dcs.installation.is_using_dcs_standalone_edition") -def test_get_dcs_install_directory_steam( - mock_standalone_edition: Mock, - mock_steam_edition: Mock, - mock_get_steam_library_folders: Mock, - tmp_path: Path, -) -> None: - install_dir = tmp_path / "steamapps/common/DCSWorld" - install_dir.mkdir(parents=True) - - mock_standalone_edition.return_value = False - mock_steam_edition.return_value = True - mock_get_steam_library_folders.return_value = ["foo", str(tmp_path), "bar"] - assert get_dcs_install_directory() == f"{install_dir}\\" +def test_get_dcs_install_directory_steam(steam_dcs_install: Path) -> None: + assert get_dcs_install_directory() == f"{steam_dcs_install}\\" -@patch("dcs.installation.is_using_dcs_steam_edition") -@patch("dcs.installation.is_using_dcs_standalone_edition") -def test_get_dcs_install_directory_not_installed( - mock_standalone_edition: Mock, - mock_steam_edition: Mock, -) -> None: - mock_standalone_edition.return_value = False - mock_steam_edition.return_value = False - +def test_get_dcs_install_directory_not_installed(none_installed: None) -> None: assert get_dcs_install_directory() == ""