From 44c4f4045f7c9cd82d7126bfc7510c80879bc08a Mon Sep 17 00:00:00 2001
From: Kristian Krsnik <git@krsnik.at>
Date: Fri, 12 Apr 2024 15:17:40 +0200
Subject: [PATCH] more error handling

---
 src/main.py  | 8 +++++++-
 src/utils.py | 1 +
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/main.py b/src/main.py
index ab870e9..9c0e24c 100644
--- a/src/main.py
+++ b/src/main.py
@@ -52,7 +52,13 @@ def test_data(api_key: str, request: Request, size: str) -> StreamingResponse:
                 detail='Invalid API Key.'
             )
 
-        size = convert_to_bytes(size)
+        try:
+            size = convert_to_bytes(size)
+        except ValueError:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail='Invalid format format for size.'
+            )
 
         if size < 0:
             raise MinSizePerRequestError
diff --git a/src/utils.py b/src/utils.py
index 0854dd7..66cf09d 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -17,6 +17,7 @@ def convert_to_bytes(size: int | str) -> int:
                 return int(float(size.removesuffix(unit)) * units[unit])
                 break
 
+        raise ValueError
 
 async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
     size_left = size