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