Enhance save confirmation modal: adjust button focus behavior and update tests for new functionality

src/hosts/tui/save_confirmation_modal.py

@@ -25,7 +25,7 @@ class SaveConfirmationModal(ModalScreen):
     
     .save-confirmation-container {
         width: 60;
-        height: 12;
+        height: 15;
         background: $surface;
         border: thick $primary;
         padding: 1;
@@ -52,6 +52,10 @@ class SaveConfirmationModal(ModalScreen):
         margin: 0 1;
         min-width: 12;
     }
+    
+    .save-confirmation-button:focus {
+        border: thick $accent;
+    }
     """
 
     BINDINGS = [
@@ -90,6 +94,60 @@ class SaveConfirmationModal(ModalScreen):
                     classes="save-confirmation-button",
                 )
 
+    def on_mount(self) -> None:
+        """Called when the modal is mounted. Set focus and ensure modal captures input."""
+        # Set focus to the modal screen itself first
+        self.focus()
+        # Then focus on the Save button
+        self.call_after_refresh(self._focus_save_button)
+
+    def _focus_save_button(self) -> None:
+        """Focus the save button after refresh."""
+        save_button = self.query_one("#save-button", Button)
+        save_button.focus()
+
+    def on_key(self, event) -> None:
+        """Handle key events, ensuring tab navigation works within the modal."""
+        if event.key == "tab":
+            # Get all buttons in order
+            buttons = [
+                self.query_one("#save-button", Button),
+                self.query_one("#discard-button", Button),
+                self.query_one("#cancel-button", Button),
+            ]
+
+            # Find currently focused button and move to next
+            for i, button in enumerate(buttons):
+                if button.has_focus:
+                    next_button = buttons[(i + 1) % len(buttons)]
+                    next_button.focus()
+                    event.prevent_default()
+                    return
+
+            # If no button has focus, focus the first one
+            buttons[0].focus()
+            event.prevent_default()
+
+        elif event.key == "shift+tab":
+            # Get all buttons in order
+            buttons = [
+                self.query_one("#save-button", Button),
+                self.query_one("#discard-button", Button),
+                self.query_one("#cancel-button", Button),
+            ]
+
+            # Find currently focused button and move to previous
+            for i, button in enumerate(buttons):
+                if button.has_focus:
+                    prev_button = buttons[(i - 1) % len(buttons)]
+                    prev_button.focus()
+                    event.prevent_default()
+                    return
+
+            # If no button has focus, focus the last one
+            buttons[-1].focus()
+            event.prevent_default()
+
     def on_button_pressed(self, event: Button.Pressed) -> None:
         """Handle button presses."""
         if event.button.id == "save-button":

tests/test_save_confirmation_modal.py

@@ -6,7 +6,7 @@ This module tests the save confirmation functionality when exiting edit entry mo
 
 import pytest
 from unittest.mock import Mock, patch
-from textual.widgets import Input, Checkbox
+from textual.widgets import Input, Checkbox, Button
 
 from hosts.main import HostsManagerApp
 from hosts.core.models import HostsFile, HostEntry
@@ -56,6 +56,29 @@ class TestSaveConfirmationModal:
 
         modal.dismiss.assert_called_once_with("cancel")
 
+    @patch.object(SaveConfirmationModal, "call_after_refresh")
+    @patch.object(SaveConfirmationModal, "focus")
+    def test_on_mount_sets_focus(self, mock_focus, mock_call_after_refresh):
+        """Test that on_mount sets focus to the modal and schedules button focus."""
+        modal = SaveConfirmationModal()
+
+        modal.on_mount()
+
+        mock_focus.assert_called_once()
+        mock_call_after_refresh.assert_called_once()
+
+    @patch.object(SaveConfirmationModal, "query_one")
+    def test_focus_save_button(self, mock_query_one):
+        """Test that _focus_save_button focuses the save button."""
+        modal = SaveConfirmationModal()
+        mock_save_button = Mock()
+        mock_query_one.return_value = mock_save_button
+
+        modal._focus_save_button()
+
+        mock_query_one.assert_called_once_with("#save-button", Button)
+        mock_save_button.focus.assert_called_once()
+
 
 class TestSaveConfirmationIntegration:
     """Test cases for save confirmation integration with the main app."""