Add configurable start timeout

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()

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,15 +1136,15 @@ 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!")
 
         container = self.start_container(
             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."""

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)
 
 

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
     ):

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(