fix(selfhost): reuse locally cached digest images during updates

scripts/update/tests/test_updater.py

@@ -79,6 +79,32 @@ class Runner:
         self.matrix_containers = {}
         self.on_stop = None
         self.real_tar = False
+        self.compose_configs = {}
+        self.local_images = set()
+        self.unavailable_images = set()
+        self.downloads = []
+
+    @staticmethod
+    def image_reference(image):
+        # Docker resolves tag@digest using the repository and digest, not the tag.
+        if '@' not in image:
+            return image
+        name, digest = image.split('@', 1)
+        repository, slash, basename = name.rpartition('/')
+        return repository + slash + basename.split(':', 1)[0] + '@' + digest
+
+    def pull_image(self, image, policy):
+        reference = self.image_reference(image)
+        if policy == 'never':
+            if reference not in self.local_images:
+                raise u.UpdateError('Image missing locally')
+            return
+        if policy == 'missing' and reference in self.local_images:
+            return
+        self.downloads.append(reference)
+        if reference in self.unavailable_images:
+            raise u.UpdateError('Registry unavailable for ' + reference)
+        self.local_images.add(reference)
 
     def run(self, args, output=None, input_file=None, **kwargs):
         self.calls.append(args)
@@ -108,6 +134,7 @@ class Runner:
                                     capture_output=True, text=True)
             if result.returncode:
                 raise u.UpdateError('Compose invalid: ' + result.stderr)
+            self.compose_configs[args[args.index('-f') + 1]] = json.loads(result.stdout)
             return result.stdout
         if 'ps' in args:
             if args[:2] == ['docker', 'ps']:
@@ -121,8 +148,17 @@ class Runner:
                     container['State']['Running'] = 'start' in args
             if 'stop' in args and self.on_stop:
                 self.on_stop()
-        if 'pull' in args and self.on_pull:
-            self.on_pull()
+        if args[:2] == ['docker', 'compose'] and any(command in args for command in ('pull', 'up', 'run')):
+            services = self.compose_configs[args[args.index('-f') + 1]]['services']
+            selected = [service for name, service in services.items() if name in args]
+            if not selected:
+                selected = list(services.values())
+            for service in selected:
+                policy = (args[args.index('--policy') + 1] if '--policy' in args else 'always') if 'pull' in args else (
+                    args[args.index('--pull') + 1] if '--pull' in args else service.get('pull_policy', 'missing'))
+                self.pull_image(service['image'], policy)
+            if 'pull' in args and self.on_pull:
+                self.on_pull()
         if any('pg_database_size' in part for part in args):
             return '1000'
         if args[0] == 'du':
@@ -493,6 +529,49 @@ class UpdateTests(unittest.TestCase):
             self.run_update()
         self.assert_untouched()
 
+    def test_cached_exact_images_work_with_unavailable_registries_and_forced_pull_overrides(self):
+        self.updater.execute('check', '3.1.0')
+        services = self.updater.candidate['services']
+        refs = {name: self.runner.image_reference(service['image']) for name, service in services.items()}
+        self.runner.local_images.update(refs.values())
+        self.runner.unavailable_images.update(ref for ref in refs.values() if '@sha256:' in ref)
+        (self.root / 'docker-compose.override.yml').write_text(
+            'services:\n' + ''.join(f'  {name}:\n    pull_policy: always\n' for name in services))
+        result = self.run_update()
+        self.assertEqual(result['status'], 'succeeded')
+        self.assertEqual(set(self.runner.downloads), {refs[name] for name in ('db', 'traefik', 'node-exporter')})
+        self.assertEqual(len(self.runner.downloads), 3)
+
+    def test_missing_digest_is_downloaded_even_if_tag_or_another_digest_is_cached(self):
+        self.updater.execute('check', '3.1.0')
+        services = self.updater.candidate['services']
+        for service in services.values():
+            image = service['image']
+            if '@' in image:
+                self.runner.local_images.add(image.split('@')[0])
+                self.runner.local_images.add(self.runner.image_reference(image).split('@')[0] + '@sha256:' + '0' * 64)
+        self.assertEqual(self.run_update()['status'], 'succeeded')
+        self.assertEqual(set(self.runner.downloads), {self.runner.image_reference(s['image']) for s in services.values()})
+
+    def test_missing_unavailable_digest_stops_before_maintenance_despite_cached_tag(self):
+        self.updater.execute('check', '3.1.0')
+        image = self.updater.candidate['services']['minio']['image']
+        self.runner.local_images.add(image.split('@')[0])
+        self.runner.unavailable_images.add(self.runner.image_reference(image))
+        with self.assertRaisesRegex(u.UpdateError, 'Registry unavailable'):
+            self.run_update()
+        self.assert_untouched()
+        self.assertFalse(self.updater.status()['recoveryRequired'])
+
+    def test_cached_mutable_tag_still_requires_a_successful_registry_refresh(self):
+        self.updater.execute('check', '3.1.0')
+        services = self.updater.candidate['services']
+        self.runner.local_images.update(self.runner.image_reference(s['image']) for s in services.values())
+        self.runner.unavailable_images.add(services['db']['image'])
+        with self.assertRaisesRegex(u.UpdateError, 'Registry unavailable'):
+            self.run_update()
+        self.assert_untouched()
+
     def test_configuration_changed_during_pull_is_detected(self):
         self.runner.on_pull = lambda: (self.root / '.env').write_bytes(self.original_env + b'\nNEW=value\n')
         with self.assertRaisesRegex(u.UpdateError, 'während'):