|
31 | 31 | import time |
32 | 32 | from datetime import timedelta |
33 | 33 | from pathlib import Path |
34 | | -from typing import Any, Callable, Optional, Union |
| 34 | +from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast |
35 | 35 | from urllib.error import HTTPError, URLError |
36 | 36 | from urllib.request import Request, urlopen |
37 | 37 |
|
| 38 | +from typing_extensions import Self |
| 39 | + |
| 40 | +from testcontainers.compose import DockerCompose |
38 | 41 | from testcontainers.core.utils import setup_logger |
39 | 42 |
|
40 | 43 | # Import base classes from waiting_utils to make them available for tests |
41 | | -from .waiting_utils import WaitStrategy, WaitStrategyTarget |
| 44 | +from testcontainers.core.waiting_utils import WaitStrategy, WaitStrategyTarget |
| 45 | + |
| 46 | +if TYPE_CHECKING: |
| 47 | + from testcontainers.core.container import DockerContainer |
42 | 48 |
|
43 | 49 | logger = setup_logger(__name__) |
44 | 50 |
|
@@ -77,22 +83,6 @@ def __init__( |
77 | 83 | self._times = times |
78 | 84 | self._predicate_streams_and = predicate_streams_and |
79 | 85 |
|
80 | | - def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "LogMessageWaitStrategy": |
81 | | - """Set the maximum time to wait for the container to be ready.""" |
82 | | - if isinstance(timeout, timedelta): |
83 | | - self._startup_timeout = int(timeout.total_seconds()) |
84 | | - else: |
85 | | - self._startup_timeout = timeout |
86 | | - return self |
87 | | - |
88 | | - def with_poll_interval(self, interval: Union[float, timedelta]) -> "LogMessageWaitStrategy": |
89 | | - """Set how frequently to check if the container is ready.""" |
90 | | - if isinstance(interval, timedelta): |
91 | | - self._poll_interval = interval.total_seconds() |
92 | | - else: |
93 | | - self._poll_interval = interval |
94 | | - return self |
95 | | - |
96 | 86 | def wait_until_ready(self, container: "WaitStrategyTarget") -> None: |
97 | 87 | """ |
98 | 88 | Wait until the specified message appears in the container logs. |
@@ -198,22 +188,6 @@ def __init__(self, port: int, path: Optional[str] = "/") -> None: |
198 | 188 | self._body: Optional[str] = None |
199 | 189 | self._insecure_tls = False |
200 | 190 |
|
201 | | - def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "HttpWaitStrategy": |
202 | | - """Set the maximum time to wait for the container to be ready.""" |
203 | | - if isinstance(timeout, timedelta): |
204 | | - self._startup_timeout = int(timeout.total_seconds()) |
205 | | - else: |
206 | | - self._startup_timeout = timeout |
207 | | - return self |
208 | | - |
209 | | - def with_poll_interval(self, interval: Union[float, timedelta]) -> "HttpWaitStrategy": |
210 | | - """Set how frequently to check if the container is ready.""" |
211 | | - if isinstance(interval, timedelta): |
212 | | - self._poll_interval = interval.total_seconds() |
213 | | - else: |
214 | | - self._poll_interval = interval |
215 | | - return self |
216 | | - |
217 | 191 | @classmethod |
218 | 192 | def from_url(cls, url: str) -> "HttpWaitStrategy": |
219 | 193 | """ |
@@ -483,22 +457,6 @@ class HealthcheckWaitStrategy(WaitStrategy): |
483 | 457 | def __init__(self) -> None: |
484 | 458 | super().__init__() |
485 | 459 |
|
486 | | - def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "HealthcheckWaitStrategy": |
487 | | - """Set the maximum time to wait for the container to be ready.""" |
488 | | - if isinstance(timeout, timedelta): |
489 | | - self._startup_timeout = int(timeout.total_seconds()) |
490 | | - else: |
491 | | - self._startup_timeout = timeout |
492 | | - return self |
493 | | - |
494 | | - def with_poll_interval(self, interval: Union[float, timedelta]) -> "HealthcheckWaitStrategy": |
495 | | - """Set how frequently to check if the container is ready.""" |
496 | | - if isinstance(interval, timedelta): |
497 | | - self._poll_interval = interval.total_seconds() |
498 | | - else: |
499 | | - self._poll_interval = interval |
500 | | - return self |
501 | | - |
502 | 460 | def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
503 | 461 | """ |
504 | 462 | Wait until the container's health check reports as healthy. |
@@ -581,22 +539,6 @@ def __init__(self, port: int) -> None: |
581 | 539 | super().__init__() |
582 | 540 | self._port = port |
583 | 541 |
|
584 | | - def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "PortWaitStrategy": |
585 | | - """Set the maximum time to wait for the container to be ready.""" |
586 | | - if isinstance(timeout, timedelta): |
587 | | - self._startup_timeout = int(timeout.total_seconds()) |
588 | | - else: |
589 | | - self._startup_timeout = timeout |
590 | | - return self |
591 | | - |
592 | | - def with_poll_interval(self, interval: Union[float, timedelta]) -> "PortWaitStrategy": |
593 | | - """Set how frequently to check if the container is ready.""" |
594 | | - if isinstance(interval, timedelta): |
595 | | - self._poll_interval = interval.total_seconds() |
596 | | - else: |
597 | | - self._poll_interval = interval |
598 | | - return self |
599 | | - |
600 | 542 | def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
601 | 543 | """ |
602 | 544 | Wait until the specified port is available for connection. |
@@ -654,22 +596,6 @@ def __init__(self, file_path: Union[str, Path]) -> None: |
654 | 596 | super().__init__() |
655 | 597 | self._file_path = Path(file_path) |
656 | 598 |
|
657 | | - def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "FileExistsWaitStrategy": |
658 | | - """Set the maximum time to wait for the container to be ready.""" |
659 | | - if isinstance(timeout, timedelta): |
660 | | - self._startup_timeout = int(timeout.total_seconds()) |
661 | | - else: |
662 | | - self._startup_timeout = timeout |
663 | | - return self |
664 | | - |
665 | | - def with_poll_interval(self, interval: Union[float, timedelta]) -> "FileExistsWaitStrategy": |
666 | | - """Set how frequently to check if the container is ready.""" |
667 | | - if isinstance(interval, timedelta): |
668 | | - self._poll_interval = interval.total_seconds() |
669 | | - else: |
670 | | - self._poll_interval = interval |
671 | | - return self |
672 | | - |
673 | 599 | def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
674 | 600 | """ |
675 | 601 | Wait until the specified file exists on the host filesystem. |
@@ -718,6 +644,65 @@ def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
718 | 644 | time.sleep(self._poll_interval) |
719 | 645 |
|
720 | 646 |
|
| 647 | +class ContainerStatusWaitStrategy(WaitStrategy): |
| 648 | + """ |
| 649 | + The possible values for the container status are: |
| 650 | + created |
| 651 | + running |
| 652 | + paused |
| 653 | + restarting |
| 654 | + exited |
| 655 | + removing |
| 656 | + dead |
| 657 | + https://docs.docker.com/reference/cli/docker/container/ls/#status |
| 658 | + """ |
| 659 | + |
| 660 | + CONTINUE_STATUSES = frozenset(("created", "restarting")) |
| 661 | + |
| 662 | + def __init__(self) -> None: |
| 663 | + super().__init__() |
| 664 | + |
| 665 | + def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
| 666 | + result = self._poll(lambda: self.running(self.get_status(container))) |
| 667 | + if not result: |
| 668 | + raise TimeoutError("container did not become running") |
| 669 | + |
| 670 | + @staticmethod |
| 671 | + def running(status: str) -> bool: |
| 672 | + if status == "running": |
| 673 | + logger.debug("status is now running") |
| 674 | + return True |
| 675 | + if status in ContainerStatusWaitStrategy.CONTINUE_STATUSES: |
| 676 | + logger.debug( |
| 677 | + "status is %s, which is valid for continuing (%s)", |
| 678 | + status, |
| 679 | + ContainerStatusWaitStrategy.CONTINUE_STATUSES, |
| 680 | + ) |
| 681 | + return False |
| 682 | + raise StopIteration(f"container status not valid for continuing: {status}") |
| 683 | + |
| 684 | + def get_status(self, container: Any) -> str: |
| 685 | + from testcontainers.core.container import DockerContainer |
| 686 | + |
| 687 | + if isinstance(container, DockerContainer): |
| 688 | + return self._get_status_tc_container(container) |
| 689 | + if isinstance(container, DockerCompose): |
| 690 | + return self._get_status_compose_container(container) |
| 691 | + raise TypeError(f"not supported operation: 'get_status' for type: {type(container)}") |
| 692 | + |
| 693 | + @staticmethod |
| 694 | + def _get_status_tc_container(container: "DockerContainer") -> str: |
| 695 | + logger.debug("fetching status of container %s", container) |
| 696 | + wrapped = container.get_wrapped_container() |
| 697 | + wrapped.reload() |
| 698 | + return cast("str", wrapped.status) |
| 699 | + |
| 700 | + @staticmethod |
| 701 | + def _get_status_compose_container(container: DockerCompose) -> str: |
| 702 | + logger.debug("fetching status of compose container %s", container) |
| 703 | + raise NotImplementedError |
| 704 | + |
| 705 | + |
721 | 706 | class CompositeWaitStrategy(WaitStrategy): |
722 | 707 | """ |
723 | 708 | Wait for multiple conditions to be satisfied in sequence. |
@@ -748,42 +733,22 @@ def __init__(self, *strategies: WaitStrategy) -> None: |
748 | 733 | super().__init__() |
749 | 734 | self._strategies = list(strategies) |
750 | 735 |
|
751 | | - def with_startup_timeout(self, timeout: Union[int, timedelta]) -> "CompositeWaitStrategy": |
752 | | - """ |
753 | | - Set the startup timeout for all contained strategies. |
754 | | -
|
755 | | - Args: |
756 | | - timeout: Maximum time to wait in seconds |
757 | | -
|
758 | | - Returns: |
759 | | - self for method chaining |
760 | | - """ |
761 | | - if isinstance(timeout, timedelta): |
762 | | - self._startup_timeout = int(timeout.total_seconds()) |
763 | | - else: |
764 | | - self._startup_timeout = timeout |
765 | | - |
766 | | - for strategy in self._strategies: |
767 | | - strategy.with_startup_timeout(timeout) |
| 736 | + def with_poll_interval(self, interval: Union[float, timedelta]) -> Self: |
| 737 | + super().with_poll_interval(interval) |
| 738 | + for _strategy in self._strategies: |
| 739 | + _strategy.with_poll_interval(interval) |
768 | 740 | return self |
769 | 741 |
|
770 | | - def with_poll_interval(self, interval: Union[float, timedelta]) -> "CompositeWaitStrategy": |
771 | | - """ |
772 | | - Set the poll interval for all contained strategies. |
773 | | -
|
774 | | - Args: |
775 | | - interval: How frequently to check in seconds |
776 | | -
|
777 | | - Returns: |
778 | | - self for method chaining |
779 | | - """ |
780 | | - if isinstance(interval, timedelta): |
781 | | - self._poll_interval = interval.total_seconds() |
782 | | - else: |
783 | | - self._poll_interval = interval |
| 742 | + def with_startup_timeout(self, timeout: Union[int, timedelta]) -> Self: |
| 743 | + super().with_startup_timeout(timeout) |
| 744 | + for _strategy in self._strategies: |
| 745 | + _strategy.with_startup_timeout(timeout) |
| 746 | + return self |
784 | 747 |
|
785 | | - for strategy in self._strategies: |
786 | | - strategy.with_poll_interval(interval) |
| 748 | + def with_transient_exceptions(self, *transient_exceptions: type[Exception]) -> Self: |
| 749 | + super().with_transient_exceptions(*transient_exceptions) |
| 750 | + for _strategy in self._strategies: |
| 751 | + _strategy.with_transient_exceptions(*transient_exceptions) |
787 | 752 | return self |
788 | 753 |
|
789 | 754 | def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
@@ -816,6 +781,7 @@ def wait_until_ready(self, container: WaitStrategyTarget) -> None: |
816 | 781 |
|
817 | 782 | __all__ = [ |
818 | 783 | "CompositeWaitStrategy", |
| 784 | + "ContainerStatusWaitStrategy", |
819 | 785 | "FileExistsWaitStrategy", |
820 | 786 | "HealthcheckWaitStrategy", |
821 | 787 | "HttpWaitStrategy", |
|
0 commit comments