control-mode: stabilize capture_pane and env assertions

Author: tony

src/libtmux/pane.py

@@ -10,6 +10,7 @@
 import dataclasses
 import logging
 import pathlib
+import time
 import typing as t
 import warnings
 
@@ -350,9 +351,25 @@ def capture_pane(
         if end is not None:
             cmd.extend(["-E", str(end)])
         output = self.cmd(*cmd).stdout
-        # Control mode (and tmux in general) can include viewport-sized blank lines.
-        while output and output[-1].strip() == "":
-            output.pop()
+
+        def _trim(lines: list[str]) -> list[str]:
+            trimmed = list(lines)
+            while trimmed and trimmed[-1].strip() == "":
+                trimmed.pop()
+            return trimmed
+
+        output = _trim(output)
+
+        # In control mode, capture-pane can race the shell: the first capture
+        # right after send-keys may return only the echoed command. Retry
+        # briefly to allow the prompt/output to land.
+        engine_name = self.server.engine.__class__.__name__
+        if engine_name == "ControlModeEngine" and not output:
+            deadline = time.monotonic() + 0.35
+            while not output and time.monotonic() < deadline:
+                time.sleep(0.05)
+                output = _trim(self.cmd(*cmd).stdout)
+
         return output
 
     def send_keys(

tests/helpers.py

@@ -0,0 +1,30 @@
+"""Test helpers for control-mode flakiness handling."""
+
+from __future__ import annotations
+
+import time
+import typing as t
+
+from libtmux.pane import Pane
+
+
+def wait_for_line(
+    pane: Pane,
+    predicate: t.Callable[[str], bool],
+    *,
+    timeout: float = 1.0,
+    interval: float = 0.05,
+) -> list[str]:
+    """Poll capture_pane until a line satisfies ``predicate``.
+
+    Returns the final capture buffer (may be empty if timeout elapses).
+    """
+    deadline = time.monotonic() + timeout
+    last: list[str] = []
+    while time.monotonic() < deadline:
+        captured = pane.capture_pane()
+        last = [captured] if isinstance(captured, str) else list(captured)
+        if any(predicate(line) for line in last):
+            break
+        time.sleep(interval)
+    return last

tests/legacy_api/test_pane.py

@@ -6,8 +6,6 @@
 import shutil
 import typing as t
 
-import pytest
-
 if t.TYPE_CHECKING:
     from libtmux.session import Session
 
@@ -77,9 +75,6 @@ def test_set_width(session: Session) -> None:
 
 def test_capture_pane(session: Session) -> None:
     """Verify Pane.capture_pane()."""
-    if session.server.engine.__class__.__name__ == "ControlModeEngine":
-        pytest.xfail("control-mode capture-pane normalization pending")
-
     env = shutil.which("env")
     assert env is not None, "Cannot find usable `env` in PATH."
 

tests/legacy_api/test_window.py

@@ -4,7 +4,6 @@
 
 import logging
 import shutil
-import time
 import typing as t
 
 import pytest
@@ -14,6 +13,7 @@
 from libtmux.pane import Pane
 from libtmux.server import Server
 from libtmux.window import Window
+from tests.helpers import wait_for_line
 
 if t.TYPE_CHECKING:
     from libtmux.session import Session
@@ -417,11 +417,14 @@ def test_split_window_with_environment(
         environment=environment,
     )
     assert pane is not None
-    # wait a bit for the prompt to be ready as the test gets flaky otherwise
-    time.sleep(0.05)
     for k, v in environment.items():
         pane.send_keys(f"echo ${k}")
-        assert pane.capture_pane()[-2] == v
+
+        def _match(line: str, expected: str = v) -> bool:
+            return line.strip() == expected
+
+        lines = wait_for_line(pane, _match)
+        assert any(_match(line) for line in lines)
 
 
 @pytest.mark.skipif(

tests/test_pane.py

@@ -67,9 +67,6 @@ def test_set_width(session: Session) -> None:
 
 def test_capture_pane(session: Session) -> None:
     """Verify Pane.capture_pane()."""
-    if session.server.engine.__class__.__name__ == "ControlModeEngine":
-        pytest.xfail("control-mode capture-pane trailing prompt normalization pending")
-
     env = shutil.which("env")
     assert env is not None, "Cannot find usable `env` in PATH."
 

tests/test_session.py

@@ -17,6 +17,7 @@
 from libtmux.test.constants import TEST_SESSION_PREFIX
 from libtmux.test.random import namer
 from libtmux.window import Window
+from tests.helpers import wait_for_line
 
 if t.TYPE_CHECKING:
     from libtmux._internal.types import StrPath
@@ -330,9 +331,6 @@ def test_new_window_with_environment(
     environment: dict[str, str],
 ) -> None:
     """Verify new window with environment vars."""
-    if session.server.engine.__class__.__name__ == "ControlModeEngine":
-        pytest.xfail("control-mode -e propagation still pending")
-
     env = shutil.which("env")
     assert env is not None, "Cannot find usable `env` in PATH."
 
@@ -346,7 +344,12 @@ def test_new_window_with_environment(
     assert pane is not None
     for k, v in environment.items():
         pane.send_keys(f"echo ${k}")
-        assert pane.capture_pane()[-2] == v
+
+        def _match(line: str, expected: str = v) -> bool:
+            return line.strip() == expected
+
+        lines = wait_for_line(pane, _match)
+        assert any(_match(line) for line in lines)
 
 
 @pytest.mark.skipif(

tests/test_window.py

@@ -5,7 +5,6 @@
 import logging
 import pathlib
 import shutil
-import time
 import typing as t
 
 import pytest
@@ -21,6 +20,7 @@
 from libtmux.pane import Pane
 from libtmux.server import Server
 from libtmux.window import Window
+from tests.helpers import wait_for_line
 
 if t.TYPE_CHECKING:
     from libtmux._internal.types import StrPath
@@ -456,11 +456,14 @@ def test_split_with_environment(
         environment=environment,
     )
     assert pane is not None
-    # wait a bit for the prompt to be ready as the test gets flaky otherwise
-    time.sleep(0.05)
     for k, v in environment.items():
         pane.send_keys(f"echo ${k}")
-        assert pane.capture_pane()[-2] == v
+
+        def _match(line: str, expected: str = v) -> bool:
+            return line.strip() == expected
+
+        lines = wait_for_line(pane, _match)
+        assert any(_match(line) for line in lines)
 
 
 @pytest.mark.skipif(