Refactor progress

Author: brunorijsman

common/exceptions.py

@@ -304,3 +304,16 @@ def __init__(self):
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             message="PSRD data is invalid.",
         )
+
+
+class InvalidEncodedFragment(DSKEException):
+    """
+    Exception raised when trying to parse an encoded fragment string that is invalid.
+    """
+
+    def __init__(self, encoded_fragment: str):
+        super().__init__(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            message="Invalid encoded fragment.",
+            details={"encoded_fragment": encoded_fragment},
+        )

common/fragment.py

@@ -6,7 +6,7 @@
 from uuid import UUID
 import pydantic
 from common.block import Block
-from common.exceptions import InvalidBlockUUIDError
+from common.exceptions import InvalidBlockUUIDError, InvalidEncodedFragment
 from . import utils
 
 
@@ -152,11 +152,13 @@ def from_enc_str(
         # TODO: Add expected_max_size parameter to avoid insane large sizes.
         parts = enc_str.split(":")
         if len(parts) != 3:
-            raise ValueError(f"Invalid fragment parameter string: {enc_str}")
+            raise InvalidEncodedFragment(encoded_fragment=enc_str)
         block_uuid_str, start_byte_str, size_str = parts
-        block = pool.get_block(UUID(block_uuid_str))
-        if block is None:
-            raise ValueError(f"Block not found: {block_uuid_str}")
+        try:
+            block_uuid = UUID(block_uuid_str)
+        except ValueError as exc:
+            raise InvalidBlockUUIDError(block_uuid=block_uuid_str) from exc
+        block = pool.get_block(block_uuid)
         start_byte = int(start_byte_str)
         size = int(size_str)
         data = block.take_data(start_byte, size)

common/tests/test_fragment.py

@@ -153,12 +153,61 @@ def test_from_api_success():
     assert block._data == bytes.fromhex("00000000000506070809")
 
 
-def test_from_api_bad_uuid():
+def test_from_api_bad_block_uuid():
     """
     Attempt to create a Fragment from a bad APIFragment (invalid block UUID).
     """
     # pylint: disable=protected-access
-    (pool, block) = _create_test_pool_and_block(10)
+    (pool, _block) = _create_test_pool_and_block(10)
+    # UUID string is not correctly formatted
     api_fragment = APIFragment(block_uuid="not-a-uuid", start=0, size=5)
     with pytest.raises(InvalidBlockUUIDError):
-        fragment = Fragment.from_api(api_fragment, pool)
+        _fragment = Fragment.from_api(api_fragment, pool)
+    # UUID string is correctly formatted, but not the UUID of any block in the pool
+    uuid = uuid4()
+    api_fragment = APIFragment(block_uuid=str(uuid), start=0, size=5)
+    with pytest.raises(InvalidBlockUUIDError):
+        _fragment = Fragment.from_api(api_fragment, pool)
+
+
+def test_to_enc_str():
+    """
+    Create an APIFragment for an encoded string.
+    """
+    block = _create_test_block(10)
+    fragment = Fragment.allocate(block, 5)
+    enc_str = fragment.to_enc_str()
+    assert enc_str == f"{block.uuid}:0:5"
+
+
+def test_from_enc_str_success():
+    """
+    Create a Fragment from a valid encoded string.
+    """
+    # pylint: disable=protected-access
+    (pool, block) = _create_test_pool_and_block(10)
+    enc_str = f"{block.uuid}:0:5"
+    fragment = Fragment.from_enc_str(enc_str, pool)
+    assert fragment.block == block
+    assert fragment.start == 0
+    assert fragment.size == 5
+    assert fragment.data == bytes.fromhex("0001020304")
+    assert block.nr_used_bytes == 5
+    assert block._data == bytes.fromhex("00000000000506070809")
+
+
+def test_from_enc_str_bad_block_uuid():
+    """
+    Attempt to create a Fragment from a bad encoded string (invalid block UUID).
+    """
+    # pylint: disable=protected-access
+    (pool, _block) = _create_test_pool_and_block(10)
+    # UUID string is not correctly formatted
+    enc_str = "not-a-uuid:0:5"
+    with pytest.raises(InvalidBlockUUIDError):
+        _fragment = Fragment.from_enc_str(enc_str, pool)
+    # UUID string is correctly formatted, but not the UUID of any block in the pool
+    uuid = uuid4()
+    enc_str = f"{uuid}:0:5"
+    with pytest.raises(InvalidBlockUUIDError):
+        _fragment = Fragment.from_enc_str(enc_str, pool)