tests: Add tests for context manager

Author: tony

tests/test_pane.py

@@ -322,3 +322,16 @@ def test_split_pane_size(session: Session) -> None:
 
         new_pane = new_pane.split(direction=PaneDirection.Right, size="10%")
         assert new_pane.pane_width == str(int(window_width_before * 0.1))
+
+
+def test_pane_context_manager(session: Session) -> None:
+    """Test Pane context manager functionality."""
+    window = session.new_window()
+    with window.split() as pane:
+        pane.send_keys('echo "Hello"')
+        assert pane in window.panes
+        assert len(window.panes) == 2  # Initial pane + new pane
+
+    # Pane should be killed after exiting context
+    assert pane not in window.panes
+    assert len(window.panes) == 1  # Only initial pane remains

tests/test_server.py

@@ -296,3 +296,15 @@ def socket_name_factory() -> str:
             myserver.kill()
         if myserver2.is_alive():
             myserver2.kill()
+
+
+def test_server_context_manager(TempServer: type[Server]) -> None:
+    """Test Server context manager functionality."""
+    with TempServer() as server:
+        session = server.new_session()
+        assert server.is_alive()
+        assert len(server.sessions) == 1
+        assert session in server.sessions
+
+    # Server should be killed after exiting context
+    assert not server.is_alive()

tests/test_session.py

@@ -381,3 +381,15 @@ def test_session_new_window_with_direction_logs_warning_for_old_tmux(
     assert any("Direction flag ignored" in record.msg for record in caplog.records), (
         "Warning missing"
     )
+
+
+def test_session_context_manager(server: Server) -> None:
+    """Test Session context manager functionality."""
+    with server.new_session() as session:
+        window = session.new_window()
+        assert session in server.sessions
+        assert window in session.windows
+        assert len(session.windows) == 2  # Initial window + new window
+
+    # Session should be killed after exiting context
+    assert session not in server.sessions

tests/test_window.py

@@ -603,3 +603,15 @@ def test_new_window_with_direction_logs_warning_for_old_tmux(
     assert any("Direction flag ignored" in record.msg for record in caplog.records), (
         "Warning missing"
     )
+
+
+def test_window_context_manager(session: Session) -> None:
+    """Test Window context manager functionality."""
+    with session.new_window() as window:
+        pane = window.split()
+        assert window in session.windows
+        assert pane in window.panes
+        assert len(window.panes) == 2  # Initial pane + new pane
+
+    # Window should be killed after exiting context
+    assert window not in session.windows