From c191f983feed280f0303460117500066731c894b Mon Sep 17 00:00:00 2001
From: Kristian Krsnik <git@krsnik.at>
Date: Fri, 5 Jul 2024 20:09:52 +0200
Subject: [PATCH] added header and small fixes

---
 flake.nix   |  2 +-
 src/main.py | 21 ++++++++++++---------
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/flake.nix b/flake.nix
index 54b7fd2..a853b7e 100644
--- a/flake.nix
+++ b/flake.nix
@@ -1,5 +1,5 @@
 {
-  description = "PLACEHOLDER";
+  description = "A webserver to create files for tetsing purposes";
 
   inputs = {
     nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
diff --git a/src/main.py b/src/main.py
index 4d44d66..fd4a2b8 100644
--- a/src/main.py
+++ b/src/main.py
@@ -44,7 +44,7 @@ class MinSizePerRequestError(Exception):
 
 
 @api.get('/')
-async def test_data(api_key: str, request: Request, size: str) -> StreamingResponse:
+async def test_data(api_key: str, size: str) -> StreamingResponse:
     try:
         if api_key not in AUTHORIZED_KEYS:
             raise HTTPException(
@@ -54,11 +54,11 @@ async def test_data(api_key: str, request: Request, size: str) -> StreamingRespo
 
         try:
             size = convert_to_bytes(size)
-        except ValueError:
+        except ValueError as err:
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,
                 detail='Invalid format format for size.'
-            )
+            ) from err
 
         if size < 0:
             raise MinSizePerRequestError
@@ -78,19 +78,22 @@ async def test_data(api_key: str, request: Request, size: str) -> StreamingRespo
         return StreamingResponse(
             status_code=status.HTTP_200_OK,
             content=generate_data(size, BUFFER_SIZE),
-            media_type='application/octet-stream'
+            media_type='application/octet-stream',
+            headers={
+                'Content-Length': size
+            }
         )
 
-    except MinSizePerRequestError:
+    except MinSizePerRequestError as err:
         raise HTTPException(
             status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
-            detail=f'Size has to be not-negative.'
-        )
-    except MaxSizePerRequestError:
+            detail='Size has to be not-negative.'
+        ) from err
+    except MaxSizePerRequestError as err:
         raise HTTPException(
             status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
             detail=f'Exceeded max size per request of {MAX_SIZE} Bytes.'
-        )
+        ) from err
 
 
 def main():