1
"""Snapshot association, interruption recovery and immutable flash inputs."""
2
from pathlib import Path
3
import tempfile
4
from types import SimpleNamespace
5
import unittest
6
 
7
from flash_cache import FlashCache, flash_images, layout_digest
8
 
9
 
10
class FlashCacheTests(unittest.TestCase):
11
    def test_diff_files_follow_sorted_addresses_and_incomplete_cache_uses_skip(self):
12
        with tempfile.TemporaryDirectory() as directory:
13
            root = Path(directory)
14
            old = root / "old.bin"
15
            old.write_bytes(b"old music")
16
            firmware = root / "firmware.bin"
17
            firmware.write_bytes(b"firmware")
18
            music = root / "music.bin"
19
            music.write_bytes(b"new music")
20
            cache = FlashCache("board-a", "layout-a", root / "cache")
21
            with cache.locked():
22
                cache.remember([(0x410000, old)])
23
                commands = []
24
                def run(command):
25
                    commands.append(command)
26
                    # A caller may prepare another catalog while this flash runs.
27
                    music.write_bytes(b"later catalog")
28
                    return SimpleNamespace(returncode=0)
29
                self.assertEqual(flash_images([(0x410000, music), (0x10000, firmware)], ["esptool"], run, cache), 0)
30
                command = commands[0]
31
                start = command.index("--diff-with")
32
                self.assertEqual(command[start + 1], "skip")
33
                self.assertEqual(Path(command[start + 2]).name, "00410000.bin")
34
                self.assertEqual(command[start + 3], "--flash-mode")
35
                self.assertNotIn("--trust-flash-content", command)
36
                self.assertNotIn("--skip-flashed", command)
37
                saved = cache.bases([0x410000])[0]
38
                self.assertEqual(saved.read_bytes(), b"new music")
39
 
40
    def test_failed_flash_retains_old_snapshot_and_full_mode_disables_fast_options(self):
41
        with tempfile.TemporaryDirectory() as directory:
42
            root = Path(directory)
43
            image = root / "image.bin"
44
            image.write_bytes(b"previous")
45
            cache = FlashCache("board", "layout", root / "cache")
46
            with cache.locked():
47
                cache.remember([(0x10000, image)])
48
                image.write_bytes(b"new")
49
                def fail(command):
50
                    self.assertNotIn("--diff-with", command)
51
                    self.assertNotIn("--skip-flashed", command)
52
                    return SimpleNamespace(returncode=2)
53
                self.assertEqual(flash_images([(0x10000, image)], ["esptool"], fail, cache, full=True), 2)
54
                self.assertEqual(cache.bases([0x10000])[0].read_bytes(), b"previous")
55
 
56
    def test_wrong_device_layout_or_corrupt_cache_cannot_supply_a_baseline(self):
57
        with tempfile.TemporaryDirectory() as directory:
58
            root = Path(directory)
59
            image = root / "image.bin"
60
            image.write_bytes(b"payload")
61
            cache = FlashCache("board-a", "layout-a", root / "cache")
62
            with cache.locked():
63
                cache.remember([(0x10000, image)])
64
                self.assertEqual(FlashCache("board-b", "layout-a", root / "cache").bases([0x10000]), [None])
65
                self.assertEqual(FlashCache("board-a", "layout-b", root / "cache").bases([0x10000]), [None])
66
                cache.bases([0x10000])[0].write_bytes(b"damaged")
67
                self.assertEqual(cache.bases([0x10000]), [None])
68
                def run(command):
69
                    self.assertIn("--skip-flashed", command)
70
                    self.assertNotIn("--diff-with", command)
71
                    return SimpleNamespace(returncode=0)
72
                flash_images([(0x10000, image)], ["esptool"], run, cache)
73
                (cache.directory / "verified.json").write_text("[]")
74
                self.assertEqual(cache.bases([0x10000]), [None])
75
 
76
    def test_sector_padding_does_not_change_partition_identity(self):
77
        table = b"partition data".ljust(0xC00, b"\xff")
78
        self.assertEqual(layout_digest(table), layout_digest(table + b"\xff" * 0x400))
79
        self.assertNotEqual(layout_digest(table), layout_digest(b"other partition"))
80
 
81
 
82
if __name__ == "__main__":
83
    unittest.main()