diff --git a/src/deadline_worker_agent/config/__main__.py b/src/deadline_worker_agent/config/__main__.py index 826821a0..7eefacc1 100644 --- a/src/deadline_worker_agent/config/__main__.py +++ b/src/deadline_worker_agent/config/__main__.py @@ -1,4 +1,5 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# type: ignore from __future__ import annotations import argparse @@ -41,6 +42,9 @@ class ParsedArguments(argparse.Namespace): """A Windows username to override the queue jobRunAs configuration. If False, then no modification is made. If None, then the setting is unset.""" + session_root_dir: str | None = None + """Path to the parent directory where the worker agent creates per-session subdirectories under""" + def create_argument_parser() -> argparse.ArgumentParser: """Creates the argparse ArgumentParser for the deadline_worker_agent.config module""" @@ -62,6 +66,12 @@ def create_argument_parser() -> argparse.ArgumentParser: help="The unique identifier for the Deadline Cloud fleet", required=False, ) + worker_group.add_argument( + "--session-root-dir", + help="The parent directory that worker agent creates per-session subdirectories under", + required=False, + type=str, + ) aws_group = parser.add_argument_group("AWS", "Settings related to AWS and EC2 hosts") parser.set_defaults(allow_ec2_instance_profile=None) @@ -193,6 +203,13 @@ def args_to_setting_modifications(parsed_args: ParsedArguments) -> list[SettingM raise NotImplementedError( f"Unexpected value for windows_job_user: {parsed_args.windows_job_user}" ) + if parsed_args.session_root_dir is not None: + settings_to_modify.append( + SettingModification( + setting=ModifiableSetting.SESSION_ROOT_DIR, + value=parsed_args.session_root_dir, + ) + ) return settings_to_modify diff --git a/src/deadline_worker_agent/config/config_file.py b/src/deadline_worker_agent/config/config_file.py index e52bbbcd..05012536 100644 --- a/src/deadline_worker_agent/config/config_file.py +++ b/src/deadline_worker_agent/config/config_file.py @@ -111,6 +111,20 @@ class ModifiableSetting(Enum): To prevent the worker agent from shutting down the host when being told to stop, uncomment the line below: +""".lstrip(), + ) + SESSION_ROOT_DIR = ModifiableSettingData( + setting_name="session_root_dir", + table_name="worker", + preceding_comment=""" +The session root directory is a parent directory where worker agent creates per-session +subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +variable is set or the --session-root-dir command-line argument is specified. + +The default session root directory on POSIX systems is "/sessions" and on Windows systems is +"C:\ProgramData\Amazon\OpenJD". + +Uncomment the line below and replace the value with your desired session root directory: """.lstrip(), ) WINDOWS_JOB_USER = ModifiableSettingData( diff --git a/src/deadline_worker_agent/installer/__init__.py b/src/deadline_worker_agent/installer/__init__.py index 52f1c424..8d739e1a 100644 --- a/src/deadline_worker_agent/installer/__init__.py +++ b/src/deadline_worker_agent/installer/__init__.py @@ -10,6 +10,11 @@ import sys import sysconfig +from deadline_worker_agent.config.settings import ( + DEFAULT_POSIX_SESSION_ROOT_DIR, + DEFAULT_WINDOWS_SESSION_ROOT_DIR, +) + if sys.platform == "win32": from deadline_worker_agent.installer.win_installer import ( @@ -70,7 +75,7 @@ def install() -> None: sys.exit(1) arg_parser = get_argument_parser() - args = arg_parser.parse_args(namespace=ParsedCommandLineArguments) + args = arg_parser.parse_args(namespace=ParsedCommandLineArguments()) scripts_path = Path(sysconfig.get_path("scripts")) if args.region is None: @@ -91,6 +96,7 @@ def install() -> None: parser=arg_parser, grant_required_access=args.grant_required_access, allow_ec2_instance_profile=not args.disallow_instance_profile, + session_root_dir=args.session_root_dir, ) if args.user: installer_args.update(user_name=args.user) @@ -124,6 +130,8 @@ def install() -> None: str(scripts_path), "--python-interpreter-path", sys.executable, + "--session-root-dir", + str(args.session_root_dir), ] if args.vfs_install_path: cmd += ["--vfs-install-path", args.vfs_install_path] @@ -169,6 +177,7 @@ class ParsedCommandLineArguments(Namespace): grant_required_access: bool disallow_instance_profile: bool windows_job_user: Optional[str] = None + session_root_dir: Path def get_argument_parser() -> ArgumentParser: # pragma: no cover @@ -259,6 +268,16 @@ def get_argument_parser() -> ArgumentParser: # pragma: no cover action="store_true", default=False, ) + parser.add_argument( + "--session-root-dir", + help="The root directory under which the worker agent creates session directories", + type=Path, + default=( + str(DEFAULT_WINDOWS_SESSION_ROOT_DIR) + if sys.platform == "win32" + else str(DEFAULT_POSIX_SESSION_ROOT_DIR) + ), # pragma: nocover + ) if sys.platform == "win32": parser.add_argument( diff --git a/src/deadline_worker_agent/installer/install.sh b/src/deadline_worker_agent/installer/install.sh index 819d1de0..e88e6a6c 100755 --- a/src/deadline_worker_agent/installer/install.sh +++ b/src/deadline_worker_agent/installer/install.sh @@ -48,6 +48,7 @@ telemetry_opt_out="no" warning_lines=() vfs_install_path="unset" python_interpreter_path=unset +session_root_dir="/sessions" usage() { @@ -106,6 +107,8 @@ usage() echo " agent makes requests to the EC2 instance meta-data service (IMDS) to check for an instance profile. " echo " If an instance profile is detected, the worker agent will stop and exit. When this is not provided, " echo " the worker agent no longer performs these checks, allowing it to run with an EC2 instance profile." + echo " --session-root-dir SESSION_ROOT_DIR" + echo " The root directory under which the worker agent will create session directories." exit 2 } @@ -131,7 +134,7 @@ validate_deadline_id() { } # Validate arguments -PARSED_ARGUMENTS=$(getopt -n install.sh --longoptions farm-id:,fleet-id:,region:,user:,group:,scripts-path:,python-interpreter-path:,vfs-install-path:,start,allow-shutdown,no-install-service,telemetry-opt-out,disallow-instance-profile -- "y" "$@") +PARSED_ARGUMENTS=$(getopt -n install.sh --longoptions farm-id:,fleet-id:,region:,user:,group:,scripts-path:,python-interpreter-path:,vfs-install-path:,session-root-dir:,start,allow-shutdown,no-install-service,telemetry-opt-out,disallow-instance-profile -- "y" "$@") VALID_ARGUMENTS=$? if [ "${VALID_ARGUMENTS}" != "0" ]; then usage @@ -152,6 +155,7 @@ do --scripts-path) scripts_path="$2" ; shift 2 ;; --python-interpreter-path) python_interpreter_path="$2" ; shift 2 ;; --vfs-install-path) vfs_install_path="$2" ; shift 2 ;; + --session-root-dir) session_root_dir="$2" ; shift 2 ;; --allow-shutdown) allow_shutdown="yes" ; shift ;; --disallow-instance-profile) disallow_instance_profile="yes" ; shift ;; --no-install-service) no_install_service="yes" ; shift ;; @@ -276,6 +280,7 @@ echo "Worker agent user: ${wa_user}" echo "Worker agent group: ${wa_group}" echo "Worker job group: ${job_group}" echo "Scripts path: ${scripts_path}" +echo "Session root directory: ${session_root_dir}" echo "Worker agent program path: ${worker_agent_program}" echo "Deadline client program path: ${client_library_program}" echo "Allow worker agent shutdown: ${allow_shutdown}" @@ -382,10 +387,12 @@ if [ -f /var/lib/deadline/worker.json ]; then fi echo "Done provisioning persistence directory (/var/lib/deadline)" -echo "Provisioning root directory for OpenJD Sessions (/sessions)" -mkdir -p /sessions -chown "${wa_user}:${job_group}" /sessions -chmod 755 /sessions +# Provision session directory +echo "Provisioning root directory for OpenJD Sessions (${session_root_dir})" +mkdir -p "${session_root_dir}" +chown "${wa_user}:${job_group}" "${session_root_dir}" +chmod 755 "${session_root_dir}" +echo "Done provisioning root directory for OpenJD Sessions (${session_root_dir})" echo "Provisioning configuration directory (/etc/amazon/deadline)" mkdir -p /etc/amazon/deadline @@ -420,7 +427,8 @@ fi --farm-id "${farm_id}" \ --fleet-id "${fleet_id}" \ "${allow_ec2_instance_profile_flag}" \ - "${shutdown_on_stop_flag}" + "${shutdown_on_stop_flag}" \ + --session-root-dir "${session_root_dir}" if ! [[ "${no_install_service}" == "yes" ]]; then # Set up the service diff --git a/src/deadline_worker_agent/installer/win_installer.py b/src/deadline_worker_agent/installer/win_installer.py index f7e172e6..6e5eac8b 100644 --- a/src/deadline_worker_agent/installer/win_installer.py +++ b/src/deadline_worker_agent/installer/win_installer.py @@ -316,6 +316,7 @@ def update_config_file( shutdown_on_stop: Optional[bool] = None, allow_ec2_instance_profile: Optional[bool] = None, windows_job_user: Optional[str] = None, + session_root_dir: Optional[Path] = None, ) -> None: """ Updates the worker configuration file, creating it from the example if it does not exist. @@ -370,6 +371,13 @@ def update_config_file( value=allow_ec2_instance_profile, ) ) + if session_root_dir is not None: + settings_to_modify.append( + SettingModification( + setting=ModifiableSetting.SESSION_ROOT_DIR, + value=str(session_root_dir), + ) + ) updated_keys = [sm.setting.value.setting_name for sm in settings_to_modify] @@ -382,7 +390,11 @@ def update_config_file( logging.info(f"Done configuring {updated_keys} in {config_path}") -def provision_directories(agent_username: str) -> WorkerAgentDirectories: +def provision_directories( + *, + agent_username: str, + session_root_dir: Path, +) -> WorkerAgentDirectories: """ Creates all required directories for Deadline Worker Agent. This function creates the following directories: @@ -394,6 +406,8 @@ def provision_directories(agent_username: str) -> WorkerAgentDirectories: Parameters agent_username(str): Worker Agent's username used for setting the permission for the directories + session_root_dir(Path): Path to the parent directory where the worker agent will create session directories + under Returns WorkerAgentDirectories: all directories created in the function @@ -433,6 +447,18 @@ def provision_directories(agent_username: str) -> WorkerAgentDirectories: os.makedirs(deadline_config_subdir, exist_ok=True) logging.info(f"Done provisioning config directory ({deadline_config_subdir})") + logging.info(f"Porvisioning session root directory ({session_root_dir})") + os.makedirs(session_root_dir, exist_ok=True) + _set_windows_permissions( + path=session_root_dir, + user=agent_username, + user_permission=FileSystemPermissionEnum.FULL_CONTROL, + group="Administrators", + group_permission=FileSystemPermissionEnum.FULL_CONTROL, + agent_user_permission=None, + ) + logging.info(f"Done provisioning session root directory ({session_root_dir})") + return WorkerAgentDirectories( deadline_dir=Path(deadline_dir), deadline_log_subdir=Path(deadline_log_subdir), @@ -775,6 +801,7 @@ def start_windows_installer( region: str, allow_shutdown: bool, parser: ArgumentParser, + session_root_dir: Path, user_name: str = DEFAULT_WA_USER, password: Optional[str] = None, group_name: str = DEFAULT_JOB_GROUP, @@ -854,6 +881,7 @@ def print_helping_info_and_exit(): f"Region: {region}\n" f"Worker agent user: {user_name}\n" f"Worker job group: {group_name}\n" + f"Session root directory: {session_root_dir}\n" f"Allow worker agent shutdown: {allow_shutdown}\n" f"Install Windows service: {install_service}\n" f"Start service: {start_service}\n" @@ -949,7 +977,7 @@ def print_helping_info_and_exit(): add_user_to_group(group_name, user_name) # Create directories and configure their permissions - agent_dirs = provision_directories(user_name) + agent_dirs = provision_directories(agent_username=user_name, session_root_dir=session_root_dir) update_config_file( str(agent_dirs.deadline_config_subdir), farm_id, @@ -959,6 +987,7 @@ def print_helping_info_and_exit(): shutdown_on_stop=allow_shutdown, allow_ec2_instance_profile=allow_ec2_instance_profile, windows_job_user=windows_job_user, + session_root_dir=session_root_dir, ) if telemetry_opt_out: diff --git a/test/integ/config/data/commented/session_root_dir/input.toml b/test/integ/config/data/commented/session_root_dir/input.toml new file mode 100644 index 00000000..06d59aad --- /dev/null +++ b/test/integ/config/data/commented/session_root_dir/input.toml @@ -0,0 +1,12 @@ +[worker] + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +# +# session_root_dir = "/sessions" diff --git a/test/integ/config/data/commented/session_root_dir/output.toml b/test/integ/config/data/commented/session_root_dir/output.toml new file mode 100644 index 00000000..fff28232 --- /dev/null +++ b/test/integ/config/data/commented/session_root_dir/output.toml @@ -0,0 +1,12 @@ +[worker] + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +# +session_root_dir = "/custom-session-root" diff --git a/test/integ/config/data/existing/session_root_dir/input.toml b/test/integ/config/data/existing/session_root_dir/input.toml new file mode 100644 index 00000000..9b2a1bf7 --- /dev/null +++ b/test/integ/config/data/existing/session_root_dir/input.toml @@ -0,0 +1,12 @@ +[worker] + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +# +session_root_dir = "/sessions" diff --git a/test/integ/config/data/existing/session_root_dir/output.toml b/test/integ/config/data/existing/session_root_dir/output.toml new file mode 100644 index 00000000..fff28232 --- /dev/null +++ b/test/integ/config/data/existing/session_root_dir/output.toml @@ -0,0 +1,12 @@ +[worker] + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +# +session_root_dir = "/custom-session-root" diff --git a/test/integ/config/data/missing/session_root_dir/input.toml b/test/integ/config/data/missing/session_root_dir/input.toml new file mode 100644 index 00000000..3e421866 --- /dev/null +++ b/test/integ/config/data/missing/session_root_dir/input.toml @@ -0,0 +1,7 @@ +[worker] + +# Some prior content that should be preserved +prior_setting_1 = 1 + +# More content to preserve +prior_setting_2 = "abc" # and here's an inline comment diff --git a/test/integ/config/data/missing/session_root_dir/output.toml b/test/integ/config/data/missing/session_root_dir/output.toml new file mode 100644 index 00000000..019d4e0b --- /dev/null +++ b/test/integ/config/data/missing/session_root_dir/output.toml @@ -0,0 +1,18 @@ +[worker] + +# Some prior content that should be preserved +prior_setting_1 = 1 + +# More content to preserve +prior_setting_2 = "abc" # and here's an inline comment + + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +session_root_dir = "/custom-session-root" diff --git a/test/integ/config/data/unset/session_root_dir/input.toml b/test/integ/config/data/unset/session_root_dir/input.toml new file mode 100644 index 00000000..fff28232 --- /dev/null +++ b/test/integ/config/data/unset/session_root_dir/input.toml @@ -0,0 +1,12 @@ +[worker] + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +# +session_root_dir = "/custom-session-root" diff --git a/test/integ/config/data/unset/session_root_dir/output.toml b/test/integ/config/data/unset/session_root_dir/output.toml new file mode 100644 index 00000000..482f6a4c --- /dev/null +++ b/test/integ/config/data/unset/session_root_dir/output.toml @@ -0,0 +1,12 @@ +[worker] + +# The session root directory is a parent directory where worker agent creates per-session +# subdirectories under. This value is overridden when the DEADLINE_WORKER_SESSION_ROOT_DIR environment +# variable is set or the --session-root-dir command-line argument is specified. +# +# The default session root directory on POSIX systems is "/sessions" and on Windows systems is +# "C:\ProgramData\Amazon\OpenJD". +# +# Uncomment the line below and replace the value with your desired session root directory: +# +# session_root_dir = "/custom-session-root" diff --git a/test/integ/config/test_config_cli.py b/test/integ/config/test_config_cli.py index c650a7f0..39ac3df9 100644 --- a/test/integ/config/test_config_cli.py +++ b/test/integ/config/test_config_cli.py @@ -69,6 +69,15 @@ def cli_args_for_shutdown_on_stop(value: str | bool | None) -> list[str]: raise NotImplementedError(f"Unexpected value: {value}") +def cli_args_for_session_root_dir(value: str | bool | None) -> list[str]: + if value is None: + return [] + elif isinstance(value, str): + return ["--session-root-dir", value] + else: + raise NotImplementedError(f"Unexpected value: {value}") + + SETTING_TO_CLI_ARGS: dict[ config_file.ModifiableSetting, Callable[[str | bool | None], list[str]] ] = { @@ -77,6 +86,7 @@ def cli_args_for_shutdown_on_stop(value: str | bool | None) -> list[str]: config_file.ModifiableSetting.FLEET_ID: cli_args_for_fleet_id, config_file.ModifiableSetting.WINDOWS_JOB_USER: cli_args_for_windows_job_user, config_file.ModifiableSetting.SHUTDOWN_ON_STOP: cli_args_for_shutdown_on_stop, + config_file.ModifiableSetting.SESSION_ROOT_DIR: cli_args_for_session_root_dir, } diff --git a/test/integ/windows/test_installer.py b/test/integ/windows/test_installer.py index 26a44a2a..aa91fc4f 100644 --- a/test/integ/windows/test_installer.py +++ b/test/integ/windows/test_installer.py @@ -5,15 +5,16 @@ assert sys.platform == "win32" -import pathlib import os import re import sys import typing import uuid +from pathlib import Path from unittest.mock import patch import pytest +from _pytest.tmpdir import TempPathFactory import win32api import win32con @@ -112,7 +113,7 @@ def test_ensure_user_profile_exists(windows_user, windows_user_password): # THEN # Verify user profile was created by checking that the home directory exists - assert pathlib.Path(f"~{windows_user}").expanduser().exists() + assert Path(f"~{windows_user}").expanduser().exists() def delete_group(group_name: str) -> None: @@ -223,7 +224,7 @@ def test_update_config_file_creates_backup(setup_example_config): assert os.path.isfile(backup_worker_config), "Backup of worker config file was not created" -def verify_least_privilege(windows_user: str, path: pathlib.Path): +def verify_least_privilege(windows_user: str, path: Path): builtin_admin_group_sid, _, _ = win32security.LookupAccountName(None, "Administrators") user_sid, _, _ = win32security.LookupAccountName(None, windows_user) sd = win32security.GetFileSecurity( @@ -259,11 +260,13 @@ def verify_least_privilege(windows_user: str, path: pathlib.Path): def test_provision_directories( windows_user: str, - tmp_path: pathlib.Path, + tmp_path: Path, + tmp_path_factory: TempPathFactory, ): # GIVEN root_dir = tmp_path / "ProgramDataTest" root_dir.mkdir() + session_root_dir: Path = tmp_path_factory.mktemp("session_root") expected_dirs = WorkerAgentDirectories( deadline_dir=root_dir / "Amazon" / "Deadline", deadline_log_subdir=root_dir / "Amazon" / "Deadline" / "Logs", @@ -285,7 +288,10 @@ def test_provision_directories( # WHEN with patch.dict(win_installer.os.environ, {"PROGRAMDATA": str(root_dir)}): - actual_dirs = provision_directories(windows_user) + actual_dirs = provision_directories( + agent_username=windows_user, + session_root_dir=session_root_dir, + ) # THEN assert actual_dirs == expected_dirs @@ -297,9 +303,11 @@ def test_provision_directories( verify_least_privilege(windows_user, actual_dirs.deadline_persistence_subdir) assert actual_dirs.deadline_config_subdir.exists() verify_least_privilege(windows_user, actual_dirs.deadline_config_subdir) + assert session_root_dir.exists() + verify_least_privilege(windows_user, session_root_dir) -def test_update_deadline_client_config(tmp_path: pathlib.Path) -> None: +def test_update_deadline_client_config(tmp_path: Path) -> None: # GIVEN deadline_client_config_path = tmp_path / "deadline_client_config" deadline_client_config_path.touch(mode=0o644, exist_ok=False) diff --git a/test/unit/conftest.py b/test/unit/conftest.py index 687d9d4a..058b83a7 100644 --- a/test/unit/conftest.py +++ b/test/unit/conftest.py @@ -140,6 +140,7 @@ def parsed_args( grant_required_access: bool, disallow_instance_profile: bool, windows_job_user: str, + session_root_dir: Path, ) -> ParsedCommandLineArguments: parsed_args = ParsedCommandLineArguments() parsed_args.farm_id = farm_id @@ -157,6 +158,7 @@ def parsed_args( parsed_args.grant_required_access = grant_required_access parsed_args.disallow_instance_profile = disallow_instance_profile parsed_args.windows_job_user = windows_job_user + parsed_args.session_root_dir = session_root_dir return parsed_args diff --git a/test/unit/install/test_install.py b/test/unit/install/test_install.py index fc890045..0bdf3351 100644 --- a/test/unit/install/test_install.py +++ b/test/unit/install/test_install.py @@ -5,7 +5,7 @@ from pathlib import Path from subprocess import CalledProcessError from typing import Generator -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import sys import sysconfig import typing @@ -67,6 +67,8 @@ def expected_cmd( sysconfig.get_path("scripts"), "--python-interpreter-path", sys.executable, + "--session-root-dir", + str(parsed_args.session_root_dir), "--vfs-install-path", parsed_args.vfs_install_path, ] @@ -142,6 +144,7 @@ class TestInstallRunsCommand: grant_required_access=True, disallow_instance_profile=True, windows_job_user="job-user", + session_root_dir="/sessions/root", ), ParsedCommandLineArguments( farm_id="farm-2", @@ -159,6 +162,7 @@ class TestInstallRunsCommand: grant_required_access=False, disallow_instance_profile=False, windows_job_user="another-job-user", + session_root_dir="/different/root", ), ) ) @@ -183,7 +187,8 @@ def test_runs_expected_subprocess( # THEN mock_subprocess_run.assert_called_once_with(expected_cmd, check=True) mock_get_arg_parser.assert_called_once_with() - mock_parse_args.assert_called_once_with(namespace=ParsedCommandLineArguments) + mock_parse_args.assert_called_once_with(namespace=ANY) + assert isinstance(mock_parse_args.call_args.kwargs["namespace"], ParsedCommandLineArguments) @pytest.mark.parametrize( diff --git a/test/unit/install/test_windows_installer.py b/test/unit/install/test_windows_installer.py index 8a370d63..702000b7 100644 --- a/test/unit/install/test_windows_installer.py +++ b/test/unit/install/test_windows_installer.py @@ -6,6 +6,8 @@ from typing import Generator from unittest.mock import Mock, call, patch, MagicMock, ANY +from deadline_worker_agent.file_system_operations import FileSystemPermissionEnum + import pytest if sys.platform != "win32": @@ -48,6 +50,7 @@ def parsed_kwargs(parsed_args: ParsedCommandLineArguments) -> Generator[dict, No "grant_required_access": parsed_args.grant_required_access, "allow_ec2_instance_profile": not parsed_args.disallow_instance_profile, "windows_job_user": parsed_args.windows_job_user, + "session_root_dir": parsed_args.session_root_dir, } @@ -950,12 +953,45 @@ def test_creates_queue_persistence_dir( agent_username: str, ) -> None: """Tests that provision_directories creates the queue persistence directory""" + # GIVEN + session_root_dir = MagicMock() # WHEN - win_installer.provision_directories(agent_username) + win_installer.provision_directories( + agent_username=agent_username, session_root_dir=session_root_dir + ) # THEN mock_os_makedirs.assert_any_call( os.path.join(os.environ["PROGRAMDATA"], "Amazon", "Deadline", "Cache", "queues"), exist_ok=True, ) + + def test_creates_session_root_dir( + self, + mock_os_makedirs: MagicMock, + mock_set_windows_permissions: MagicMock, + agent_username: str, + ) -> None: + """Tests that provision_directories creates the session root directory""" + # GIVEN + session_root_dir = MagicMock() + + # WHEN + win_installer.provision_directories( + agent_username=agent_username, session_root_dir=session_root_dir + ) + + # THEN + mock_os_makedirs.assert_any_call( + session_root_dir, + exist_ok=True, + ) + mock_set_windows_permissions.assert_called_with( + path=session_root_dir, + user=agent_username, + user_permission=FileSystemPermissionEnum.FULL_CONTROL, + group="Administrators", + group_permission=FileSystemPermissionEnum.FULL_CONTROL, + agent_user_permission=None, + )