From e35d8dbf3ddc9e65a4120aa7ca79988ab5f6f0fb Mon Sep 17 00:00:00 2001
From: "Brian S. Stephan" <bss@incorporeal.org>
Date: Fri, 12 Apr 2024 16:58:53 -0500
Subject: [PATCH] add some more UF2 tests and sanity checks

Signed-off-by: Brian S. Stephan <bss@incorporeal.org>
---
 gp2040ce_bintools/storage.py |  4 ++++
 tests/test_builder.py        |  1 -
 tests/test_storage.py        | 21 +++++++++++++++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/gp2040ce_bintools/storage.py b/gp2040ce_bintools/storage.py
index 9fea627..8e877c7 100644
--- a/gp2040ce_bintools/storage.py
+++ b/gp2040ce_bintools/storage.py
@@ -117,6 +117,10 @@ def convert_uf2_to_binary(uf2: bytearray) -> bytearray:
 
         binary += content[0:bytes_]
         old_uf2_addr = uf2_addr
+
+    # when this is all done we should have counted the expected number of blocks
+    if block_count != block_num + 1:
+        raise ValueError(f"not all expected blocks ({block_count}) were found, only got {block_num + 1}!")
     return binary
 
 
diff --git a/tests/test_builder.py b/tests/test_builder.py
index 59b5385..d406f68 100644
--- a/tests/test_builder.py
+++ b/tests/test_builder.py
@@ -3,7 +3,6 @@
 SPDX-FileCopyrightText: © 2023 Brian S. Stephan <bss@incorporeal.org>
 SPDX-License-Identifier: GPL-3.0-or-later
 """
-import math
 import os
 import sys
 import unittest.mock as mock
diff --git a/tests/test_storage.py b/tests/test_storage.py
index 17cd934..4c9a363 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -166,6 +166,27 @@ def test_convert_binary_to_uf2_to_binary(whole_board_with_board_config_dump):
     assert whole_board_with_board_config_dump == binary
 
 
+def test_malformed_uf2(whole_board_with_board_config_dump):
+    """Check that we expect a properly-formed UF2."""
+    uf2 = storage.convert_binary_to_uf2(whole_board_with_board_config_dump)
+
+    # truncated UF2 --- byte mismatch
+    with pytest.raises(ValueError):
+        storage.convert_uf2_to_binary(uf2[:-4])
+
+    # truncated uf2 --- counter is wrong
+    with pytest.raises(ValueError):
+        storage.convert_uf2_to_binary(uf2[512:])
+
+    # truncated uf2 --- total count is wrong
+    with pytest.raises(ValueError):
+        storage.convert_uf2_to_binary(uf2[:-512])
+
+    # malformed UF2 --- counter jumps in the middle, suggests total blocks is wrong
+    with pytest.raises(ValueError):
+        storage.convert_uf2_to_binary(uf2 + uf2)
+
+
 @with_pb2s
 def test_serialize_config_with_footer(storage_dump, config_binary):
     """Test that reserializing a read in config matches the original.