|
6 | 6 | - HealthcheckWaitStrategy: Wait for Docker health checks to pass |
7 | 7 | - PortWaitStrategy: Wait for TCP ports to be available |
8 | 8 | - FileExistsWaitStrategy: Wait for files to exist on the filesystem |
| 9 | +- ExecWaitStrategy: Wait for command execution inside container to succeed |
9 | 10 | - CompositeWaitStrategy: Combine multiple wait strategies |
10 | 11 |
|
11 | 12 | Example: |
|
19 | 20 | # Wait for log message |
20 | 21 | container.waiting_for(LogMessageWaitStrategy("Server started")) |
21 | 22 |
|
| 23 | + # Wait for command execution |
| 24 | + container.waiting_for(ExecWaitStrategy(["pg_isready", "-U", "postgres"])) |
| 25 | +
|
22 | 26 | # Combine multiple strategies |
23 | 27 | container.waiting_for(CompositeWaitStrategy( |
24 | 28 | LogMessageWaitStrategy("Database ready"), |
@@ -779,9 +783,103 @@ def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
779 | 783 | logger.debug("CompositeWaitStrategy: All strategies completed successfully") |
780 | 784 |
|
781 | 785 |
|
| 786 | +class ExecWaitStrategy(WaitStrategy): |
| 787 | + """ |
| 788 | + Wait for a command execution inside the container to succeed. |
| 789 | +
|
| 790 | + This strategy executes a command inside the container and waits for it to |
| 791 | + return a successful exit code. It's useful for databases and services |
| 792 | + that provide CLI tools to check readiness. |
| 793 | +
|
| 794 | + Args: |
| 795 | + command: Command to execute (list of strings or single string) |
| 796 | + expected_exit_code: Expected exit code for success (default: 0) |
| 797 | +
|
| 798 | + Example: |
| 799 | + # Wait for Postgres readiness |
| 800 | + strategy = ExecWaitStrategy( |
| 801 | + ["sh", "-c", |
| 802 | + "PGPASSWORD='password' psql -U user -d db -h 127.0.0.1 -c 'select 1;'"] |
| 803 | + ) |
| 804 | +
|
| 805 | + # Wait for Redis readiness |
| 806 | + strategy = ExecWaitStrategy(["redis-cli", "ping"]) |
| 807 | +
|
| 808 | + # Check for specific exit code |
| 809 | + strategy = ExecWaitStrategy(["custom-healthcheck.sh"], expected_exit_code=0) |
| 810 | + """ |
| 811 | + |
| 812 | + def __init__( |
| 813 | + self, |
| 814 | + command: Union[str, list[str]], |
| 815 | + expected_exit_code: int = 0, |
| 816 | + ) -> None: |
| 817 | + super().__init__() |
| 818 | + self._command = command if isinstance(command, list) else [command] |
| 819 | + self._expected_exit_code = expected_exit_code |
| 820 | + |
| 821 | + def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
| 822 | + """ |
| 823 | + Wait until command execution succeeds with the expected exit code. |
| 824 | +
|
| 825 | + Args: |
| 826 | + container: The container to execute commands in |
| 827 | +
|
| 828 | + Raises: |
| 829 | + TimeoutError: If the command doesn't succeed within the timeout period |
| 830 | + RuntimeError: If the container doesn't support exec |
| 831 | + """ |
| 832 | + # Check if container supports exec (DockerContainer does, ComposeContainer doesn't) |
| 833 | + if not hasattr(container, "exec"): |
| 834 | + raise RuntimeError( |
| 835 | + f"ExecWaitStrategy requires a container with exec support. " |
| 836 | + f"Container type {type(container).__name__} does not support exec." |
| 837 | + ) |
| 838 | + |
| 839 | + start_time = time.time() |
| 840 | + last_exit_code = None |
| 841 | + last_output = None |
| 842 | + |
| 843 | + while True: |
| 844 | + duration = time.time() - start_time |
| 845 | + if duration > self._startup_timeout: |
| 846 | + command_str = " ".join(self._command) |
| 847 | + raise TimeoutError( |
| 848 | + f"Command execution did not succeed within {self._startup_timeout:.3f} seconds. " |
| 849 | + f"Command: {command_str}. " |
| 850 | + f"Expected exit code: {self._expected_exit_code}, " |
| 851 | + f"last exit code: {last_exit_code}. " |
| 852 | + f"Last output: {last_output}. " |
| 853 | + f"Hint: Check if the service is starting correctly, the command is valid, " |
| 854 | + f"and all required environment variables or credentials are properly configured." |
| 855 | + ) |
| 856 | + |
| 857 | + try: |
| 858 | + result = container.exec(self._command) |
| 859 | + last_exit_code = result.exit_code |
| 860 | + last_output = result.output.decode() if hasattr(result.output, "decode") else str(result.output) |
| 861 | + |
| 862 | + if result.exit_code == self._expected_exit_code: |
| 863 | + logger.debug( |
| 864 | + f"ExecWaitStrategy: Command succeeded with exit code {result.exit_code} after {duration:.2f}s" |
| 865 | + ) |
| 866 | + return |
| 867 | + |
| 868 | + logger.debug( |
| 869 | + f"ExecWaitStrategy: Command failed with exit code {result.exit_code}, " |
| 870 | + f"expected {self._expected_exit_code}. Retrying..." |
| 871 | + ) |
| 872 | + except Exception as e: |
| 873 | + logger.debug(f"ExecWaitStrategy: Command execution failed with exception: {e}. Retrying...") |
| 874 | + last_output = str(e) |
| 875 | + |
| 876 | + time.sleep(self._poll_interval) |
| 877 | + |
| 878 | + |
782 | 879 | __all__ = [ |
783 | 880 | "CompositeWaitStrategy", |
784 | 881 | "ContainerStatusWaitStrategy", |
| 882 | + "ExecWaitStrategy", |
785 | 883 | "FileExistsWaitStrategy", |
786 | 884 | "HealthcheckWaitStrategy", |
787 | 885 | "HttpWaitStrategy", |
|
0 commit comments