add IDAllocation.reserve method

Allow reserving a particular range instead of allocating a suitable
large range of an IDAllocation. This is useful when a directory
hierarchy defines the allocation and we merely want to verify it to be
assigned.

1 files changed, 18 insertions, 0 deletions

diff --git a/tests/test_simple.py b/tests/test_simple.py
index 212f414..456e088 100644
--- a/tests/test_simple.py
+++ b/tests/test_simple.py
@@ -48,6 +48,24 @@ class IDAllocationTest(unittest.TestCase):
         alloc.add_range(3, 2)
         self.assertIn(alloc.allocate(3), (1, 2))
 
+    def test_reserve(self) -> None:
+        alloc = linuxnamespaces.IDAllocation()
+        alloc.add_range(0, 10)
+        # Split a range
+        alloc.reserve(3, 3)
+        self.assertEqual(alloc.ranges, [(0, 3), (6, 4)])
+        self.assertRaises(ValueError, alloc.reserve, 0, 4)
+        self.assertRaises(ValueError, alloc.reserve, 5, 4)
+        # Head of range
+        alloc.reserve(0, 2)
+        self.assertEqual(alloc.ranges, [(2, 1), (6, 4)])
+        # Tail of range
+        alloc.reserve(7, 3)
+        self.assertEqual(alloc.ranges, [(2, 1), (6, 1)])
+        # Exact range
+        alloc.reserve(6, 1)
+        self.assertEqual(alloc.ranges, [(2, 1)])
+
 
 class AsnycioTest(unittest.IsolatedAsyncioTestCase):
     async def test_eventfd(self) -> None: