probert: implement os prober arg

subiquity/cmd/server.py

@@ -73,6 +73,8 @@ def make_server_args_parser():
         help='in dryrun, control basedir of files')
     parser.add_argument(
         '--storage-version', action='store', type=int, default=1)
+    parser.add_argument(
+        '--use-os-prober', action='store_true', default=False)
     return parser
 
 

subiquity/server/controllers/filesystem.py

@@ -408,7 +408,9 @@ class FilesystemController(SubiquityController, FilesystemManipulator):
             fname = 'probe-data-restricted.json'
             key = "ProbeDataRestricted"
         else:
-            probe_types = None
+            probe_types = {'defaults'}
+            if self.app.opts.use_os_prober:
+                probe_types |= {'os'}
             fname = 'probe-data.json'
             key = "ProbeData"
         storage = await run_in_thread(

subiquity/server/controllers/tests/test_filesystem.py

@@ -0,0 +1,46 @@
+# Copyright 2022 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 unittest import mock, TestCase
+
+from subiquity.server.controllers.filesystem import FilesystemController
+
+from subiquitycore.tests.util import run_coro
+from subiquitycore.tests.mocks import make_app
+
+
+class TestSubiquityControllerFilesystem(TestCase):
+    def setUp(self):
+        self.app = make_app()
+        self.app.opts.bootloader = 'UEFI'
+        self.app.report_start_event = mock.Mock()
+        self.app.report_finish_event = mock.Mock()
+        self.app.prober = mock.Mock()
+        self.fsc = FilesystemController(app=self.app)
+        self.fsc._configured = True
+
+    def test_probe_restricted(self):
+        run_coro(self.fsc._probe_once(context=None, restricted=True))
+        self.app.prober.get_storage.assert_called_with({'blockdev'})
+
+    def test_probe_defaults(self):
+        self.app.opts.use_os_prober = False
+        run_coro(self.fsc._probe_once(context=None, restricted=False))
+        self.app.prober.get_storage.assert_called_with({'defaults'})
+
+    def test_probe_defaults_and_os(self):
+        self.app.opts.use_os_prober = True
+        run_coro(self.fsc._probe_once(context=None, restricted=False))
+        self.app.prober.get_storage.assert_called_with({'defaults', 'os'})

subiquitycore/prober.py

@@ -50,7 +50,7 @@ class Prober():
                 time.sleep(2)
                 1/0
             r = self.saved_config['storage'].copy()
-            if probe_types is not None:
+            if probe_types is not None and 'defaults' not in probe_types:
                 for k in self.saved_config['storage']:
                     if k not in probe_types:
                         r[k] = {}

subiquitycore/tests/test_prober.py

@@ -0,0 +1,27 @@
+# Copyright 2022 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.tests import SubiTestCase
+
+from subiquitycore.prober import Prober
+
+
+class TestProber(SubiTestCase):
+    def test_none_and_defaults_equal(self):
+        with open('examples/simple.json', 'r') as fp:
+            prober = Prober(machine_config=fp, debug_flags=())
+        none_storage = prober.get_storage(probe_types=None)
+        defaults_storage = prober.get_storage(probe_types={'defaults'})
+        self.assertEqual(defaults_storage, none_storage)