ui: have a distinct state file for rich mode over serial

We recently made sure that after doing a snap refresh, the rich mode
(i.e., either rich or basic) is preserved. This was implemented by
storing the rich mode in a state file. When the client starts, it loads
the rich mode from said state file if it exists.

Unfortunately, on s390x, it causes installs to default to basic mode.
This happens because on this architecture, a subiquity install consists
of:

 * a first client (over serial) showing the SSH password
 * a second client (logging over SSH) actually going through the
   installation UI.

Since the first client uses a serial connection, the state file is
created with rich-mode set to basic. Upon connecting using SSH, the
state file is read and the rich-mode is set to basic as well.

Fixed by storing the rich-mode in two separate files, one for clients
over serial and one for other clients.

LP: #2036096

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
(cherry picked from commit c95261e0de)

subiquitycore/tests/test_tui.py

@@ -0,0 +1,76 @@
+# Copyright 2023 Canonical, Ltd.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import json
+import pathlib
+from unittest.mock import Mock, patch
+
+from subiquitycore.tests import SubiTestCase
+from subiquitycore.tui import TuiApplication
+
+
+class TestTuiApplication(SubiTestCase):
+    def setUp(self):
+        with patch("subiquitycore.tui.Application.__init__", return_value=None):
+            opts = Mock()
+            opts.answers = None
+            opts.dry_run = True
+            self.tui = TuiApplication(opts)
+            # Usually, the below are assigned by Application.__init__()
+            self.tui.opts = opts
+            self.tui.state_dir = self.tmp_dir()
+
+    def test_get_initial_rich_mode_normal(self):
+        self.tui.opts.run_on_serial = False
+        self.assertTrue(self.tui.get_initial_rich_mode())
+
+        # With a state file.
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode-tty").open("w") as fh:
+            fh.write(json.dumps(True))
+        self.assertTrue(self.tui.get_initial_rich_mode())
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode-tty").open("w") as fh:
+            fh.write(json.dumps(False))
+        self.assertFalse(self.tui.get_initial_rich_mode())
+
+    def test_get_initial_rich_mode_serial(self):
+        self.tui.opts.run_on_serial = True
+        self.assertFalse(self.tui.get_initial_rich_mode())
+
+        # With a state file.
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode-serial").open("w") as fh:
+            fh.write(json.dumps(True))
+        self.assertTrue(self.tui.get_initial_rich_mode())
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode-serial").open("w") as fh:
+            fh.write(json.dumps(False))
+        self.assertFalse(self.tui.get_initial_rich_mode())
+
+    def test_get_initial_rich_mode_legacy_state_file(self):
+        # Make sure if an old rich-mode state file is present, it is honored -
+        # but the new format takes precedence.
+        self.tui.opts.run_on_serial = True
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode").open("w") as fh:
+            fh.write(json.dumps(True))
+        self.assertTrue(self.tui.get_initial_rich_mode())
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode-serial").open("w") as fh:
+            fh.write(json.dumps(False))
+        self.assertFalse(self.tui.get_initial_rich_mode())
+
+        self.tui.opts.run_on_serial = False
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode").open("w") as fh:
+            fh.write(json.dumps(False))
+        self.assertFalse(self.tui.get_initial_rich_mode())
+        with (pathlib.Path(self.tui.state_dir) / "rich-mode-tty").open("w") as fh:
+            fh.write(json.dumps(True))
+        self.assertTrue(self.tui.get_initial_rich_mode())

subiquitycore/tui.py

@@ -263,7 +263,11 @@ class TuiApplication(Application):
             urwid.util.set_encoding("ascii")
             new_palette = PALETTE_MONO
             self.rich_mode = False
-        with open(self.state_path("rich-mode"), "w") as fp:
+        if self.opts.run_on_serial:
+            rich_mode_file = "rich-mode-serial"
+        else:
+            rich_mode_file = "rich-mode-tty"
+        with open(self.state_path(rich_mode_file), "w") as fp:
             json.dump(self.rich_mode, fp)
         urwid.CanvasCache.clear()
         self.urwid_loop.screen.register_palette(new_palette)
@@ -286,6 +290,37 @@ class TuiApplication(Application):
     def make_screen(self, inputf=None, outputf=None):
         return make_screen(self.opts.ascii, inputf, outputf)
 
+    def get_initial_rich_mode(self) -> bool:
+        """Return the initial value for rich-mode, either loaded from an
+        exising state file or automatically determined. True means rich mode
+        and False means basic mode."""
+        if self.opts.run_on_serial:
+            rich_mode_file = "rich-mode-serial"
+        else:
+            rich_mode_file = "rich-mode-tty"
+        try:
+            fp = open(self.state_path(rich_mode_file))
+        except FileNotFoundError:
+            pass
+        else:
+            with fp:
+                return json.load(fp)
+
+        try:
+            # During the 23.10 development cycle, there was only one rich-mode
+            # state file. Let's handle the scenario where we just snap
+            # refresh-ed from a pre 23.10 release.
+            # Once mantic is EOL, let's remove this code.
+            fp = open(self.state_path("rich-mode"))
+        except FileNotFoundError:
+            pass
+        else:
+            with fp:
+                return json.load(fp)
+
+        # By default, basic on serial - rich otherwise.
+        return not self.opts.run_on_serial
+
     def start_urwid(self, input=None, output=None):
         # This stops the tcsetpgrp call in run_command_in_foreground from
         # suspending us. See the rant there for more details.
@@ -302,14 +337,7 @@ class TuiApplication(Application):
             **self.extra_urwid_loop_args(),
         )
         extend_dec_special_charmap()
-        try:
-            fp = open(self.state_path("rich-mode"))
-        except FileNotFoundError:
-            initial_rich_mode = not self.opts.run_on_serial
-        else:
-            with fp:
-                initial_rich_mode = json.load(fp)
-        self.set_rich(initial_rich_mode)
+        self.set_rich(self.get_initial_rich_mode())
         self.urwid_loop.start()
         self.select_initial_screen()