From 9a11e9071939db74dd83dccdd5403eaea3c33de9 Mon Sep 17 00:00:00 2001
From: Michael Hudson-Doyle <michael.hudson@canonical.com>
Date: Mon, 6 Apr 2020 16:28:46 +1200
Subject: [PATCH] attach a controller to most contexts

---
 subiquity/context.py    | 26 ++++++++++++++++++++++++++
 subiquity/controller.py |  1 +
 subiquity/core.py       |  3 +++
 subiquitycore/core.py   |  3 ++-
 4 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 subiquity/context.py

diff --git a/subiquity/context.py b/subiquity/context.py
new file mode 100644
index 00000000..6f294c7e
--- /dev/null
+++ b/subiquity/context.py
@@ -0,0 +1,26 @@
+# Copyright 2020 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/>.
+
+from subiquitycore.context import Context
+
+
+class SubiquityContext(Context):
+
+    controller = None
+
+    def __init__(self, app, name, description, parent, level, childlevel=None):
+        super().__init__(app, name, description, parent, level, childlevel)
+        if parent is not None:
+            self.controller = parent.controller
diff --git a/subiquity/controller.py b/subiquity/controller.py
index 5f102a93..304e5d03 100644
--- a/subiquity/controller.py
+++ b/subiquity/controller.py
@@ -35,6 +35,7 @@ class SubiquityController(BaseController):
     def __init__(self, app):
         super().__init__(app)
         self.autoinstall_applied = False
+        self.context.controller = self
         self.setup_autoinstall()
 
     def setup_autoinstall(self):
diff --git a/subiquity/core.py b/subiquity/core.py
index 7ffbd382..9b769047 100644
--- a/subiquity/core.py
+++ b/subiquity/core.py
@@ -37,6 +37,7 @@ from subiquitycore.controller import Skip
 from subiquitycore.core import Application
 from subiquitycore.utils import run_command
 
+from subiquity.context import SubiquityContext
 from subiquity.controllers.error import (
     ErrorReportKind,
     )
@@ -86,6 +87,8 @@ class Subiquity(Application):
 
     project = "subiquity"
 
+    context_cls = SubiquityContext
+
     def make_model(self):
         root = '/'
         if self.opts.dry_run:
diff --git a/subiquitycore/core.py b/subiquitycore/core.py
index 88f0b21a..ba1fec21 100644
--- a/subiquitycore/core.py
+++ b/subiquitycore/core.py
@@ -315,6 +315,7 @@ class Application:
     # instance.
 
     make_ui = SubiquityCoreUI
+    context_cls = Context
 
     def __init__(self, opts):
         self.debug_flags = ()
@@ -366,7 +367,7 @@ class Application:
         self.new_event_loop()
         self.urwid_loop = None
         self.controllers = ControllerSet(self, self.controllers)
-        self.context = Context.new(self)
+        self.context = self.context_cls.new(self)
 
     def new_event_loop(self):
         new_loop = asyncio.new_event_loop()