diff --git a/composer_local_dev/cli.py b/composer_local_dev/cli.py index bc8cbda9..5da362a5 100644 --- a/composer_local_dev/cli.py +++ b/composer_local_dev/cli.py @@ -179,6 +179,19 @@ def cli(): ) +option_start_timeout = click.option( + "--start-timeout", + "start_timeout_seconds", + type=click.IntRange(min=0), + help=( + "Maximum number of seconds to wait for environment startup. " + "Use 0 to disable the timeout." + ), + show_default=f"{constants.OPERATION_TIMEOUT_SECONDS} seconds", + metavar="SECONDS", +) + + def _complete_environment(ctx, param, incomplete): env_dirs = files.get_environment_directories() return [ @@ -364,12 +377,14 @@ def create( @cli.command() @optional_environment @option_port +@option_start_timeout @verbose_mode @debug_mode @errors.catch_exceptions() def start( environment: Optional[str], web_server_port: Optional[int], + start_timeout_seconds: Optional[int], verbose: bool, debug: bool, ): @@ -380,7 +395,7 @@ def start( env_path, web_server_port ) console.get_console().print(f"Starting {env.name} composer environment...") - env.start() + env.start(timeout_seconds=start_timeout_seconds) @cli.command() @@ -404,12 +419,14 @@ def stop(environment: Optional[str], verbose: bool, debug: bool): @cli.command() @optional_environment @option_port +@option_start_timeout @verbose_mode @debug_mode @errors.catch_exceptions() def restart( environment: Optional[str], web_server_port: Optional[int], + start_timeout_seconds: Optional[int], verbose: bool, debug: bool, ): @@ -424,7 +441,7 @@ def restart( env = composer_environment.Environment.load_from_config( env_path, web_server_port ) - env.restart() + env.restart(timeout_seconds=start_timeout_seconds) @cli.command() diff --git a/composer_local_dev/environment.py b/composer_local_dev/environment.py index 6326d29a..fc436404 100644 --- a/composer_local_dev/environment.py +++ b/composer_local_dev/environment.py @@ -39,9 +39,13 @@ DOCKER_FILES = pathlib.Path(__file__).parent / "docker_files" -def timeout_occurred(start_time): +def timeout_occurred(start_time, timeout_seconds=None): """Returns whether time since start is greater than OPERATION_TIMEOUT.""" - return time.time() - start_time >= constants.OPERATION_TIMEOUT_SECONDS + if timeout_seconds is None: + timeout_seconds = constants.OPERATION_TIMEOUT_SECONDS + if timeout_seconds <= 0: + return False + return time.time() - start_time >= timeout_seconds def get_image_mounts( @@ -996,7 +1000,7 @@ def assert_container_is_active(self, container_name): ): raise errors.EnvironmentStartError() - def wait_for_db_start(self): + def wait_for_db_start(self, timeout_seconds=None): start_time = time.time() with console.get_console().status("[bold green]Starting database..."): self.assert_container_is_active(self.db_container_name) @@ -1011,12 +1015,14 @@ def wait_for_db_start(self): "Database is started in %.2f seconds", start_duration ) return - if timeout_occurred(start_time): - raise errors.EnvironmentStartTimeoutError() + if timeout_occurred(start_time, timeout_seconds): + raise errors.EnvironmentStartTimeoutError( + timeout_seconds or constants.OPERATION_TIMEOUT_SECONDS + ) self.assert_container_is_active(self.db_container_name) raise errors.EnvironmentStartError() - def wait_for_start(self): + def wait_for_start(self, timeout_seconds=None): """ Poll environment logs to see if it is ready. When Airflow scheduler starts, it prints 'searching for files' in the @@ -1039,8 +1045,10 @@ def wait_for_start(self): "Environment started in %.2f seconds", start_duration ) return - if timeout_occurred(start_time): - raise errors.EnvironmentStartTimeoutError() + if timeout_occurred(start_time, timeout_seconds): + raise errors.EnvironmentStartTimeoutError( + timeout_seconds or constants.OPERATION_TIMEOUT_SECONDS + ) self.assert_container_is_active(self.container_name) raise errors.EnvironmentStartError() @@ -1090,7 +1098,7 @@ def start_container( error = f"Environment ({container_name}) failed to start with an error: {err}" raise errors.EnvironmentStartError(error) from None - def start(self, assert_not_running=True): + def start(self, assert_not_running=True, timeout_seconds=None): """Starts local composer environment. Before starting we are asserting that are required files in the @@ -1128,7 +1136,7 @@ def start(self, assert_not_running=True): f"Database engine is selected as {self.database_engine}. The container will start before" ) db_container = self.start_container(self.db_container_name, False) - self.wait_for_db_start() + self.wait_for_db_start(timeout_seconds) self.ensure_container_is_attached_to_network(db_container) LOG.info(f"Database started!") @@ -1136,7 +1144,7 @@ def start(self, assert_not_running=True): self.container_name, assert_not_running ) self.ensure_container_is_attached_to_network(container) - self.wait_for_start() + self.wait_for_start(timeout_seconds) self.print_start_message() def ensure_container_is_attached_to_network(self, container): @@ -1217,7 +1225,7 @@ def stop(self, remove_container=False): network = self.get_docker_network() network.remove() - def restart(self): + def restart(self, timeout_seconds=None): """ Restarts the local composer environment. @@ -1228,7 +1236,7 @@ def restart(self): self.stop(remove_container=True) except errors.EnvironmentNotRunningError: pass - self.start(assert_not_running=False) + self.start(assert_not_running=False, timeout_seconds=timeout_seconds) def status(self) -> str: """Get status of the local composer environment.""" diff --git a/composer_local_dev/errors.py b/composer_local_dev/errors.py index 0e899cb5..8115ceb5 100644 --- a/composer_local_dev/errors.py +++ b/composer_local_dev/errors.py @@ -95,10 +95,8 @@ def __init__(self, msg: Optional[str] = None): class EnvironmentStartTimeoutError(EnvironmentStartError): """Composer environment start timed out.""" - def __init__(self): - msg = constants.ENV_DID_NOT_START_TIMEOUT_ERROR.format( - seconds=constants.OPERATION_TIMEOUT_SECONDS - ) + def __init__(self, seconds: int = constants.OPERATION_TIMEOUT_SECONDS): + msg = constants.ENV_DID_NOT_START_TIMEOUT_ERROR.format(seconds=seconds) super().__init__(msg) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index c40fee68..53938682 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -185,12 +185,19 @@ def mocked_resolve_env(self, env_path): def assert_env_loaded(self, mocked_env, env_path, port=None): mocked_env.load_from_config.assert_called_with(env_path, port) - def assert_run_command(self, command, mocked_env, env_path, port=None): + def assert_run_command( + self, command, mocked_env, env_path, port=None, start_timeout=None + ): run_composer_and_assert_exit_code( command, exit_code=0, ) self.assert_env_loaded(mocked_env, env_path, port) + loaded_env = mocked_env.load_from_config.return_value + if command.startswith("start"): + loaded_env.start.assert_called_with(timeout_seconds=start_timeout) + elif command.startswith("restart"): + loaded_env.restart.assert_called_with(timeout_seconds=start_timeout) @pytest.mark.parametrize("command", ["start", "restart"]) def test_start_command( @@ -206,6 +213,39 @@ def test_start_command_with_port( command += f" --port {port}" self.assert_run_command(command, mocked_env, env_path, port) + @pytest.mark.parametrize("command", ["start", "restart"]) + def test_start_command_with_timeout( + self, mocked_env, mocked_resolve_env, env_path, command + ): + start_timeout = 600 + command += f" --start-timeout {start_timeout}" + self.assert_run_command( + command, mocked_env, env_path, start_timeout=start_timeout + ) + + @pytest.mark.parametrize("command", ["start", "restart"]) + def test_start_command_with_timeout_disabled( + self, mocked_env, mocked_resolve_env, env_path, command + ): + start_timeout = 0 + command += f" --start-timeout {start_timeout}" + self.assert_run_command( + command, mocked_env, env_path, start_timeout=start_timeout + ) + + @pytest.mark.parametrize("command", ["start", "restart"]) + def test_start_command_with_invalid_timeout( + self, mocked_env, mocked_resolve_env, env_path, command + ): + result = run_composer_and_assert_exit_code( + f"{command} --start-timeout -1", + exit_code=2, + ) + assert ( + "Invalid value for '--start-timeout': -1 is not in the range" + in result.output + ) + def test_start_with_invalid_port( self, mocked_env, mocked_resolve_env, env_path ): diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py index ea771d8b..06247028 100644 --- a/tests/unit/test_environment.py +++ b/tests/unit/test_environment.py @@ -1250,6 +1250,30 @@ def test_wait_for_start_timeout(self, mocked_time, default_env): ): default_env.wait_for_start() + @mock.patch( + "composer_local_dev.environment.time.time", + side_effect=[1, 121], + ) + def test_wait_for_start_custom_timeout(self, mocked_time, default_env): + container = mock.Mock() + container.status = "running" + container.logs = mock.Mock(return_value=[b"Log lines"]) + default_env.get_container = mock.Mock(return_value=container) + with pytest.raises( + errors.ComposerCliError, + match="Environment did not start in 120 seconds.", + ): + default_env.wait_for_start(timeout_seconds=120) + + @mock.patch( + "composer_local_dev.environment.time.time", + side_effect=[1, 1 + constants.OPERATION_TIMEOUT_SECONDS], + ) + def test_wait_for_start_timeout_disabled(self, mocked_time, default_env): + log_lines = [b"Log lines", b"Searching for files in path..."] + default_env.get_container = get_container_logs_mock(log_lines) + default_env.wait_for_start(timeout_seconds=0) + def test_wait_for_start_failed(self, default_env): default_env.get_container = get_container_logs_mock([], "not_running") with pytest.raises(