From d74a506f98d5e734f5b1e616c3bcabd18dc247d0 Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Mon, 28 Nov 2022 14:32:29 +0100
Subject: [PATCH] ubuntu-pro: add ability to run magic-attach locally in
 dry-run mode

In dry-run mode, the 'magic_attach_run_locally' variable can now be used
to execute uaclient on the host (against the uacontracts environment
specified in /etc/ubuntu-advantage/uaclient.conf) instead of relying on
a mock mechanism.

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 examples/dr-config-magic-attach.yaml                  | 1 +
 subiquity/cmd/schema.py                               | 4 ++++
 subiquity/server/controllers/tests/test_ubuntu_pro.py | 5 ++++-
 subiquity/server/controllers/ubuntu_pro.py            | 7 ++++++-
 subiquity/server/dryrun.py                            | 6 ++++++
 5 files changed, 21 insertions(+), 2 deletions(-)
 create mode 100644 examples/dr-config-magic-attach.yaml

diff --git a/examples/dr-config-magic-attach.yaml b/examples/dr-config-magic-attach.yaml
new file mode 100644
index 00000000..aa75f5b1
--- /dev/null
+++ b/examples/dr-config-magic-attach.yaml
@@ -0,0 +1 @@
+pro_magic_attach_run_locally: true
diff --git a/subiquity/cmd/schema.py b/subiquity/cmd/schema.py
index 7ebb646a..e12b4819 100644
--- a/subiquity/cmd/schema.py
+++ b/subiquity/cmd/schema.py
@@ -22,6 +22,7 @@ import sys
 import jsonschema
 
 from subiquity.cmd.server import make_server_args_parser
+from subiquity.server.dryrun import DRConfig
 from subiquity.server.server import SubiquityServer
 
 
@@ -54,6 +55,9 @@ def make_app():
     parser = make_server_args_parser()
     opts, unknown = parser.parse_known_args(['--dry-run'])
     app = SubiquityServer(opts, '')
+    # This is needed because the ubuntu-pro server controller accesses dr_cfg
+    # in the initializer.
+    app.dr_cfg = DRConfig()
     app.base_model = app.make_model()
     app.controllers.load_all()
     return app
diff --git a/subiquity/server/controllers/tests/test_ubuntu_pro.py b/subiquity/server/controllers/tests/test_ubuntu_pro.py
index fb03b894..99a014c8 100644
--- a/subiquity/server/controllers/tests/test_ubuntu_pro.py
+++ b/subiquity/server/controllers/tests/test_ubuntu_pro.py
@@ -18,12 +18,15 @@ import unittest
 from subiquity.server.controllers.ubuntu_pro import (
     UbuntuProController,
 )
+from subiquity.server.dryrun import DRConfig
 from subiquitycore.tests.mocks import make_app
 
 
 class TestUbuntuProController(unittest.TestCase):
     def setUp(self):
-        self.controller = UbuntuProController(make_app())
+        app = make_app()
+        app.dr_cfg = DRConfig()
+        self.controller = UbuntuProController(app)
 
     def test_serialize(self):
         self.controller.model.token = "1a2b3C"
diff --git a/subiquity/server/controllers/ubuntu_pro.py b/subiquity/server/controllers/ubuntu_pro.py
index 831ae635..3d68f16a 100644
--- a/subiquity/server/controllers/ubuntu_pro.py
+++ b/subiquity/server/controllers/ubuntu_pro.py
@@ -90,7 +90,12 @@ class UbuntuProController(SubiquityController):
         """ Initializer for server-side Ubuntu Pro controller. """
         strategy: UAInterfaceStrategy
         if app.opts.dry_run:
-            strategy = MockedUAInterfaceStrategy(scale_factor=app.scale_factor)
+            if app.dr_cfg.pro_magic_attach_run_locally:
+                executable = "/usr/bin/ubuntu-advantage"
+                strategy = UAClientUAInterfaceStrategy(executable=executable)
+            else:
+                strategy = MockedUAInterfaceStrategy(
+                        scale_factor=app.scale_factor)
         else:
             # Make sure we execute `$PYTHON "$SNAP/usr/bin/ubuntu-advantage"`.
             executable = (
diff --git a/subiquity/server/dryrun.py b/subiquity/server/dryrun.py
index 863cbdba..17c88ac0 100644
--- a/subiquity/server/dryrun.py
+++ b/subiquity/server/dryrun.py
@@ -36,6 +36,12 @@ class DRConfig:
 
     # Tells whether "$source"/var/lib/snapd/seed/systems exists on the source.
     systems_dir_exists: bool = False
+    # Tells whether we should run /usr/bin/ubuntu-advantage instead of using
+    # Mock objects.
+    pro_magic_attach_run_locally: bool = False
+    # When running /usr/bin/ubuntu-advantage locally, do not use the production
+    # ua-contrats.
+    pro_ua_contracts_url: str = "https://contracts.staging.canonical.com"
 
     @classmethod
     def load(cls, stream):