Added a test case

hirak99 · Nov 5, 2023 · 9dd7753 · 9dd7753
1 parent 12bc3b3
commit 9dd7753
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/code/snap_operator.py b/src/code/snap_operator.py
@@ -57,6 +57,7 @@ def find_target(config: configs.Config, suffix: str) -> Optional[snap_holder.Sna
 
 
 def _all_but_last_k(array: list[_GenericT], k: int) -> Iterator[_GenericT]:
+    """All but at most k last elements."""
     if k < 0:
         raise ValueError(f"k = {k} < 0")
     yield from array[: len(array) - k]

diff --git a/src/code/snap_operator_test.py b/src/code/snap_operator_test.py
@@ -140,8 +140,11 @@ def test_all_but_k(self):
         self.assertEqual(
             list(snap_operator._all_but_last_k([1, 2, 3, 4], 0)), [1, 2, 3, 4]
         )
+        # Leave -2 out of 4 elements - causes an error.
         with self.assertRaisesRegex(ValueError, r"k = .+ < 0"):
             list(snap_operator._all_but_last_k([1, 2, 3, 4], -2))
+        # Leave 5 out of 2 elements. Succeeds.
+        self.assertEqual(list(snap_operator._all_but_last_k([1, 2], 5)), [])
 
     def setUp(self) -> None:
         super().setUp()