!squsah wip

Author: tony

src/libtmux/test/waiter.py

@@ -6,18 +6,34 @@
 
 from __future__ import annotations
 
+import time
 import typing as t
 from dataclasses import dataclass
-from typing import Callable, TypeVar
+from typing import (
+    TypeVar,
+)
 
-from libtmux.test.retry import retry_until
+from libtmux.exc import LibTmuxException
+from libtmux.test.retry import WaitTimeout, retry_until
 
 if t.TYPE_CHECKING:
     from libtmux.pane import Pane
 
 T = TypeVar("T")
 
 
+class WaiterError(LibTmuxException):
+    """Base exception for waiter errors."""
+
+
+class WaiterTimeoutError(WaiterError):
+    """Exception raised when waiting for content times out."""
+
+
+class WaiterContentError(WaiterError):
+    """Exception raised when there's an error getting or checking content."""
+
+
 @dataclass
 class WaitResult(t.Generic[T]):
     """Result of a wait operation."""
@@ -43,112 +59,101 @@ def __init__(self, pane: Pane, timeout: float = 2.0) -> None:
         self.pane = pane
         self.timeout = timeout
 
-    def wait_for_content(
+    def _check_content(
         self,
-        predicate: Callable[[str], bool],
-        *,
-        timeout: float | None = None,
-        error_message: str | None = None,
-    ) -> WaitResult[str]:
-        """Wait for pane content to match predicate.
+        predicate: t.Callable[[str], bool],
+        result: WaitResult,
+    ) -> bool:
+        """Check pane content against predicate.
 
         Parameters
         ----------
         predicate : Callable[[str], bool]
             Function that takes pane content as string and returns bool
-        timeout : float | None, optional
-            Timeout in seconds, by default None (uses instance timeout)
-        error_message : str | None, optional
-            Custom error message if timeout occurs, by default None
+        result : WaitResult
+            Result object to store content if predicate matches
 
         Returns
         -------
-        WaitResult[str]
-            Result containing success status and pane content if successful
+        bool
+            True if predicate matches, False otherwise
+
+        Raises
+        ------
+        WaiterContentError
+            If there's an error capturing pane content
         """
-        timeout = timeout or self.timeout
-        result = WaitResult[str](success=False)
-
-        def check_content() -> bool:
-            try:
-                content = "\n".join(self.pane.capture_pane())
-                if predicate(content):
-                    result.success = True
-                    result.value = content
-                    return True
-                else:
-                    return False
-            except Exception as e:
-                result.error = e
-                return False
+        try:
+            content = "\n".join(self.pane.capture_pane())
+            if predicate(content):
+                result.value = content
+                return True
+            return False
+        except Exception as e:
+            error = WaiterContentError("Error capturing pane content")
+            error.__cause__ = e
+            raise error from e
 
+    def wait_for_content(
+        self,
+        predicate: t.Callable[[str], bool],
+        timeout_seconds: float | None = None,
+        interval_seconds: float | None = None,
+        error_message: str | None = None,
+    ) -> WaitResult:
+        """Wait for content in the pane to match a predicate."""
+        result = WaitResult(success=False, value=None, error=None)
         try:
-            success = retry_until(check_content, timeout, raises=False)
+            # Give the shell a moment to be ready
+            time.sleep(0.1)
+            success = retry_until(
+                lambda: self._check_content(predicate, result),
+                seconds=timeout_seconds or self.timeout,
+                interval=interval_seconds,
+                raises=True,
+            )
+            result.success = success
             if not success:
-                result.error = Exception(
+                result.error = WaiterTimeoutError(
                     error_message or "Timed out waiting for content",
                 )
-        except Exception as e:
+        except WaitTimeout as e:
+            result.error = WaiterTimeoutError(error_message or str(e))
+            result.success = False
+        except WaiterContentError as e:
             result.error = e
-            if error_message:
-                result.error = Exception(error_message)
-
+            result.success = False
+        except Exception as e:
+            if isinstance(e, (WaiterTimeoutError, WaiterContentError)):
+                result.error = e
+            else:
+                result.error = WaiterContentError("Error capturing pane content")
+                result.error.__cause__ = e
+            result.success = False
         return result
 
     def wait_for_prompt(
         self,
         prompt: str,
-        *,
-        timeout: float | None = None,
+        timeout_seconds: float | None = None,
         error_message: str | None = None,
-    ) -> WaitResult[str]:
-        """Wait for specific prompt to appear in pane.
-
-        Parameters
-        ----------
-        prompt : str
-            The prompt text to wait for
-        timeout : float | None, optional
-            Timeout in seconds, by default None (uses instance timeout)
-        error_message : str | None, optional
-            Custom error message if timeout occurs, by default None
-
-        Returns
-        -------
-        WaitResult[str]
-            Result containing success status and pane content if successful
-        """
+    ) -> WaitResult:
+        """Wait for a specific prompt to appear in the pane."""
         return self.wait_for_content(
             lambda content: prompt in content and len(content.strip()) > 0,
-            timeout=timeout,
+            timeout_seconds=timeout_seconds,
             error_message=error_message or f"Prompt '{prompt}' not found in pane",
         )
 
     def wait_for_text(
         self,
         text: str,
-        *,
-        timeout: float | None = None,
+        timeout_seconds: float | None = None,
         error_message: str | None = None,
-    ) -> WaitResult[str]:
-        """Wait for specific text to appear in pane.
-
-        Parameters
-        ----------
-        text : str
-            The text to wait for
-        timeout : float | None, optional
-            Timeout in seconds, by default None (uses instance timeout)
-        error_message : str | None, optional
-            Custom error message if timeout occurs, by default None
-
-        Returns
-        -------
-        WaitResult[str]
-            Result containing success status and pane content if successful
-        """
+    ) -> WaitResult:
+        """Wait for specific text to appear in the pane."""
         return self.wait_for_content(
             lambda content: text in content,
-            timeout=timeout,
+            timeout_seconds=timeout_seconds,
             error_message=error_message or f"Text '{text}' not found in pane",
         )

tests/test/test_waiter.py

@@ -5,9 +5,16 @@
 import shutil
 import typing as t
 
-from libtmux.test.waiter import PaneWaiter
+from libtmux.test.retry import WaitTimeout
+from libtmux.test.waiter import (
+    PaneWaiter,
+    WaiterContentError,
+    WaiterTimeoutError,
+)
 
 if t.TYPE_CHECKING:
+    from pytest import MonkeyPatch
+
     from libtmux.session import Session
 
 
@@ -72,6 +79,8 @@ def test_wait_timeout(session: Session) -> None:
     assert not result.success
     assert result.value is None
     assert result.error is not None
+    assert isinstance(result.error, WaiterTimeoutError)
+    assert str(result.error) == "Text 'this text will never appear' not found in pane"
 
 
 def test_custom_error_message(session: Session) -> None:
@@ -96,6 +105,7 @@ def test_custom_error_message(session: Session) -> None:
     assert not result.success
     assert result.value is None
     assert result.error is not None
+    assert isinstance(result.error, WaiterTimeoutError)
     assert str(result.error) == custom_message
 
 
@@ -120,3 +130,128 @@ def test_wait_for_content_predicate(session: Session) -> None:
     assert result.success
     assert result.value is not None
     assert "123" in result.value
+
+
+def test_wait_for_content_inner_exception(
+    session: Session,
+    monkeypatch: MonkeyPatch,
+) -> None:
+    """Test exception handling in wait_for_content's inner try-except."""
+    env = shutil.which("env")
+    assert env is not None, "Cannot find usable `env` in PATH."
+
+    session.new_window(
+        attach=True,
+        window_name="test_waiter",
+        window_shell=f"{env} PROMPT_COMMAND='' PS1='READY>' sh",
+    )
+    pane = session.active_window.active_pane
+    assert pane is not None
+
+    waiter = PaneWaiter(pane)
+
+    def mock_capture_pane(*args: t.Any, **kwargs: t.Any) -> list[str]:
+        """Mock capture_pane that raises an exception."""
+        msg = "Test error"
+        raise Exception(msg)
+
+    monkeypatch.setattr(pane, "capture_pane", mock_capture_pane)
+    result = waiter.wait_for_text("some text")
+    assert not result.success
+    assert result.value is None
+    assert result.error is not None
+    assert isinstance(result.error, WaiterContentError)
+    assert str(result.error) == "Error capturing pane content"
+    assert isinstance(result.error.__cause__, Exception)
+    assert str(result.error.__cause__) == "Test error"
+
+
+def test_wait_for_content_outer_exception(
+    session: Session,
+    monkeypatch: MonkeyPatch,
+) -> None:
+    """Test exception handling in wait_for_content's outer try-except."""
+    env = shutil.which("env")
+    assert env is not None, "Cannot find usable `env` in PATH."
+
+    session.new_window(
+        attach=True,
+        window_name="test_waiter",
+        window_shell=f"{env} PROMPT_COMMAND='' PS1='READY>' sh",
+    )
+    pane = session.active_window.active_pane
+    assert pane is not None
+
+    waiter = PaneWaiter(pane)
+
+    def mock_retry_until(*args: t.Any, **kwargs: t.Any) -> bool:
+        """Mock retry_until that raises an exception."""
+        msg = "Custom error"
+        raise WaitTimeout(msg)
+
+    monkeypatch.setattr("libtmux.test.waiter.retry_until", mock_retry_until)
+    result = waiter.wait_for_text(
+        "some text",
+        error_message="Custom error",
+    )
+    assert not result.success
+    assert result.value is None
+    assert result.error is not None
+    assert isinstance(result.error, WaiterTimeoutError)
+    assert str(result.error) == "Custom error"
+
+
+def test_wait_for_content_outer_exception_no_custom_message(
+    session: Session,
+    monkeypatch: MonkeyPatch,
+) -> None:
+    """Test exception handling in wait_for_content's outer try-except without custom message."""
+    env = shutil.which("env")
+    assert env is not None, "Cannot find usable `env` in PATH."
+
+    session.new_window(
+        attach=True,
+        window_name="test_waiter",
+        window_shell=f"{env} PROMPT_COMMAND='' PS1='READY>' sh",
+    )
+    pane = session.active_window.active_pane
+    assert pane is not None
+
+    waiter = PaneWaiter(pane)
+
+    def mock_capture_pane(*args: t.Any, **kwargs: t.Any) -> list[str]:
+        """Mock capture_pane that raises an exception."""
+        msg = "Test error"
+        raise Exception(msg)
+
+    monkeypatch.setattr(pane, "capture_pane", mock_capture_pane)
+    result = waiter.wait_for_text("some text")  # No custom error message
+    assert not result.success
+    assert result.value is None
+    assert result.error is not None
+    assert isinstance(result.error, WaiterContentError)
+    assert str(result.error) == "Error capturing pane content"
+    assert isinstance(result.error.__cause__, Exception)
+    assert str(result.error.__cause__) == "Test error"
+
+
+def test_wait_for_content_retry_exception(monkeypatch, session) -> None:
+    """Test that retry exceptions are handled correctly."""
+    pane = session.new_window("test_waiter").active_pane
+
+    def mock_retry_until(
+        predicate,
+        timeout_seconds=None,
+        interval_seconds=None,
+        raises=None,
+    ) -> t.NoReturn:
+        msg = "Text 'some text' not found in pane"
+        raise WaitTimeout(msg)
+
+    monkeypatch.setattr("libtmux.test.waiter.retry_until", mock_retry_until)
+    waiter = PaneWaiter(pane)
+    result = waiter.wait_for_content(lambda content: "some text" in content)
+
+    assert not result.success
+    assert result.value is None
+    assert str(result.error) == "Text 'some text' not found in pane"