From 93f06eeb0f46424329d01154ebc2f1b11249f729 Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Thu, 24 Nov 2022 17:31:17 +0100
Subject: [PATCH] dry-run: add configuration object to control dry-run behavior

Oftentimes, we want to simulate a specific behavior of the application
when running in dry-run mode. To do so, we use either command line
parameters or environment variables.

This patch introduces a configuration object for dry-run executions
only. The object can be automatically loaded from a JSON file specified
via the --dry-run-config CLI argument.

Such a configuration object should help us cover way more test cases.
Going forward, I would like to use this object for things like:

 * drivers - to instruct Subiquity what third-party drivers it should
   suggest ; or if Subiquity should run ubuntu-drivers on the host
   instead.
 * ubuntu-pro - to specify the ua-contracts test environment URL - or
   predefined automatic replies for the server
 * to assume that /var/lib/snapd/seed/systems directory exists on the
   source (or not).
 * to specify the Ubuntu release that is returned by lsb_release ; can
   be used to test behavior on LTS vs non LTS releases.
 *
 * ...

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 scripts/runtests.sh        |  4 ++++
 subiquity/cmd/server.py    | 10 ++++++++++
 subiquity/server/dryrun.py | 16 ++++++++++++++++
 subiquity/server/server.py |  2 ++
 4 files changed, 32 insertions(+)

diff --git a/scripts/runtests.sh b/scripts/runtests.sh
index e4563de8..b597b8ac 100755
--- a/scripts/runtests.sh
+++ b/scripts/runtests.sh
@@ -159,6 +159,7 @@ for answers in examples/answers*.yaml; do
     if echo $answers|grep -vq system-setup; then
         config=$(sed -n 's/^#machine-config: \(.*\)/\1/p' $answers || true)
         catalog=$(sed -n 's/^#source-catalog: \(.*\)/\1/p' $answers || true)
+        dr_config=$(sed -n 's/^#dr-config: \(.*\)/\1/p' "$answers" || true)
         if [ -z "$config" ]; then
             config=examples/simple.json
         fi
@@ -170,6 +171,9 @@ for answers in examples/answers*.yaml; do
         if [ -n "$serial" ]; then
             opts+=(--serial)
         fi
+        if [ -n "$dr_config" ]; then
+            opts+=(--dry-run-config "$dr_config")
+        fi
         # The --foreground is important to avoid subiquity getting SIGTTOU-ed.
         LANG=C.UTF-8 timeout --foreground 60 \
             python3 -m subiquity.cmd.tui < "$tty" \
diff --git a/subiquity/cmd/server.py b/subiquity/cmd/server.py
index 113952e3..6fd6e693 100644
--- a/subiquity/cmd/server.py
+++ b/subiquity/cmd/server.py
@@ -62,6 +62,7 @@ def make_server_args_parser():
     parser.add_argument('--dry-run', action='store_true',
                         dest='dry_run',
                         help='menu-only, do not call installer function')
+    parser.add_argument('--dry-run-config', type=argparse.FileType())
     parser.add_argument('--socket')
     parser.add_argument('--machine-config', metavar='CONFIG',
                         dest='machine_config',
@@ -118,6 +119,7 @@ def main():
     # setup_environment sets $APPORT_DATA_DIR which must be set before
     # apport is imported, which is done by this import:
     from subiquity.server.server import SubiquityServer
+    from subiquity.server.dryrun import DRConfig
     parser = make_server_args_parser()
     opts = parser.parse_args(sys.argv[1:])
     if opts.storage_version is None:
@@ -125,9 +127,16 @@ def main():
             'subiquity-storage-version', 1))
     logdir = LOGDIR
     if opts.dry_run:
+        if opts.dry_run_config:
+            dr_cfg = DRConfig.load(opts.dry_run_config)
+        else:
+            dr_cfg = DRConfig()
+
         if opts.snaps_from_examples is None:
             opts.snaps_from_examples = True
         logdir = opts.output_base
+    else:
+        dr_cfg = None
     if opts.socket is None:
         if opts.dry_run:
             opts.socket = opts.output_base + '/socket'
@@ -157,6 +166,7 @@ def main():
 
     async def run_with_loop():
         server = SubiquityServer(opts, block_log_dir)
+        server.dr_cfg = dr_cfg
         server.note_file_for_apport(
             "InstallerServerLog", logfiles['debug'])
         server.note_file_for_apport(
diff --git a/subiquity/server/dryrun.py b/subiquity/server/dryrun.py
index 7da5e809..1b6c58c5 100644
--- a/subiquity/server/dryrun.py
+++ b/subiquity/server/dryrun.py
@@ -13,6 +13,10 @@
 # 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 yaml
+
+import attr
+
 
 class DryRunController:
 
@@ -22,3 +26,15 @@ class DryRunController:
 
     async def crash_GET(self) -> None:
         1/0
+
+
+@attr.s(auto_attribs=True)
+class DRConfig:
+    """ Configuration for dry-run-only executions.
+    All variables here should have default values ; to indicate the behavior we
+    want by default in dry-run mode. """
+
+    @classmethod
+    def load(cls, stream):
+        data = yaml.safe_load(stream)
+        return cls(**data)
diff --git a/subiquity/server/server.py b/subiquity/server/server.py
index 70ba36c7..7841bbec 100644
--- a/subiquity/server/server.py
+++ b/subiquity/server/server.py
@@ -68,6 +68,7 @@ from subiquity.models.subiquity import (
     ModelNames,
     SubiquityModel,
     )
+from subiquity.server.dryrun import DRConfig
 from subiquity.server.controller import SubiquityController
 from subiquity.server.geoip import (
     GeoIP,
@@ -274,6 +275,7 @@ class SubiquityServer(Application):
 
     def __init__(self, opts, block_log_dir):
         super().__init__(opts)
+        self.dr_cfg: Optional[DRConfig] = None
         self.set_source_variant(self.supported_variants[0])
         self.block_log_dir = block_log_dir
         self.cloud = None