Merge pull request #1324 from ogayot/remove_overlay_no_overlay_ok

ui: avoid crashing when removing overlay that does not exist

subiquity/ui/views/ssh.py

@@ -297,5 +297,7 @@ class SSHView(BaseView):
             ConfirmSSHKeys(self, ssh_data, identities))
 
     def fetching_ssh_keys_failed(self, msg, stderr):
-        self.remove_overlay()
+        # FIXME in answers-based runs, the overlay does not exist so we pass
+        # not_found_ok=True.
+        self.remove_overlay(not_found_ok=True)
         self.show_stretchy_overlay(SomethingFailed(self, msg, stderr))

subiquitycore/tests/test_view.py

@@ -16,7 +16,7 @@
 import urwid
 
 from subiquitycore.tests import SubiTestCase
-from subiquitycore.view import BaseView
+from subiquitycore.view import BaseView, OverlayNotFoundError
 from subiquitycore.ui.stretchy import Stretchy, StretchyOverlay
 from subiquitycore.ui.utils import undisabled
 
@@ -93,3 +93,20 @@ class TestBaseView(SubiTestCase):
         bv.remove_overlay()
         self.assertTrue(c.was_closed)
         self.assertEqual(self.get_stretchy_chain(bv), [b, a])
+
+    def test_remove_overlay_not_found(self):
+        bv, a, b, c = self.make_view_with_overlays()
+        bv.remove_overlay(not_found_ok=False)
+        bv.remove_overlay(not_found_ok=False)
+        bv.remove_overlay(not_found_ok=False)
+
+        # At this point, there is no more overlay.
+        with self.assertRaises(OverlayNotFoundError):
+            bv.remove_overlay(not_found_ok=False)
+
+        bv.remove_overlay(not_found_ok=True)
+
+        with self.assertRaises(OverlayNotFoundError):
+            bv.remove_overlay(stretchy=a, not_found_ok=False)
+
+        bv.remove_overlay(stretchy=a, not_found_ok=True)

subiquitycore/view.py

@@ -38,6 +38,10 @@ from subiquitycore.ui.utils import disabled, undisabled
 log = logging.getLogger('subiquitycore.view')
 
 
+class OverlayNotFoundError(Exception):
+    """ Exception to raise when trying to remove a non-existent overlay. """
+
+
 class BaseView(WidgetWrap):
 
     def local_help(self):
@@ -78,7 +82,9 @@ class BaseView(WidgetWrap):
         stretchy.opened()
         self._w = StretchyOverlay(disabled(self._w), stretchy)
 
-    def remove_overlay(self, stretchy=None):
+    def remove_overlay(self, stretchy=None,
+                       *, not_found_ok=False) -> None:
+        """ Remove (frontmost) overlay from the view. """
         if stretchy is not None:
             one_above = None
             cur = self._w
@@ -95,10 +101,20 @@ class BaseView(WidgetWrap):
                 one_above = cur
                 cur = undisabled(cur.bottom_w)
             else:
+                if not not_found_ok:
+                    raise OverlayNotFoundError
+        else:
+            try:
+                behind_overlay = self._w.bottom_w
+            except AttributeError:
+                if not_found_ok:
+                    return
+                raise OverlayNotFoundError
+
             if isinstance(self._w, StretchyOverlay):
                 emit_signal(self._w.stretchy, 'closed')
                 self._w.stretchy.closed()
-            self._w = undisabled(self._w.bottom_w)
+            self._w = undisabled(behind_overlay)
 
     def cancel(self):
         pass
@@ -106,10 +122,9 @@ class BaseView(WidgetWrap):
     def keypress(self, size, key):
         key = super().keypress(size, key)
         if key == 'esc':
-            if hasattr(self._w, 'bottom_w'):
-                self.remove_overlay()
-                return None
-            else:
+            try:
+                self.remove_overlay(not_found_ok=False)
+            except OverlayNotFoundError:
                 self.cancel()
                 return None
         return key