updated

* Project can now be started as a package
* Added unit-tests

.gitignore

@@ -2,5 +2,13 @@
 
 /.direnv/
 
+# config and database
 *.json
-*.bin
+
+# python
+__pycache__/
+.pytest_cache/
+.mypy_cache/
+
+# FastHTML
+/.sesskey

README.md

@@ -10,7 +10,7 @@
     "log": "-",
     "buffer-size": "4KiB",
     "max-size": "2GB",
-    "keys": [
+    "api-keys": [
         "TESTKEY"
     ],
     "max-data": "10GB",

flake.nix

@@ -18,12 +18,35 @@
     forAllSystems = inputs.nixpkgs.lib.genAttrs supportedSystems;
     pkgs = forAllSystems (system: inputs.nixpkgs.legacyPackages.${system});
     poetry2nix = forAllSystems (system: inputs.poetry2nix-lib.lib.mkPoetry2Nix {pkgs = pkgs.${system};});
+
+    addSetuptools = final: prev: list:
+      builtins.listToAttrs (builtins.map (package: {
+          name = "${package}";
+          value = prev."${package}".overridePythonAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or []) ++ [final.setuptools];
+          });
+        })
+        list);
   in {
     # `nix build`
-    packages = forAllSystems (system: {
-      default = poetry2nix.${system}.mkPoetryApplication {
-        projectDir = self;
+    packages = forAllSystems (system: let
+      mkPackage = {debug ? false}:
+        poetry2nix.${system}.mkPoetryApplication {
+          projectDir = ./.;
+          checkPhase =
+            if debug
+            then "pyright --warnings testdata && pytest"
+            else "";
+          # doCheck = debug;
+          preferWheels = false;
+          nativeBuildInputs = with pkgs.${system}; [pyright];
+          overrides =
+            poetry2nix.${system}.overrides.withDefaults (final: prev:
+              addSetuptools final prev ["sqlite-minutils" "fastlite" "python-fasthtml"]);
         };
+    in {
+      default = mkPackage {debug = false;};
+      debug = mkPackage {debug = true;};
       vm = self.nixosConfigurations.vm.config.microvm.declaredRunner;
     });
 
@@ -40,18 +63,23 @@
 
     # `nix develop`
     devShells = forAllSystems (system: {
-      default = let
-        poetryEnv =
-          if builtins.pathExists ./poetry.lock
-          then poetry2nix.${system}.mkPoetryEnv {projectDir = self;}
-          else null;
-      in
-        pkgs.${system}.mkShellNoCC {
-          packages = with pkgs.${system};
-            [
-              poetry
-            ]
-            ++ [poetryEnv];
+      # Shell for app dependencies.
+      #
+      #     nix develop
+      #
+      # Use this shell for developing your app.
+      default = pkgs.${system}.mkShellNoCC {
+        inputsFrom = [self.packages.${system}.default];
+        packages = [self.packages.${system}.default];
+      };
+
+      # Shell for poetry.
+      #
+      #     nix develop .#poetry
+      #
+      # Use this shell for changes to pyproject.toml and poetry.lock.
+      poetry = pkgs.${system}.mkShellNoCC {
+        packages = [pkgs.${system}.poetry];
       };
     });
 

pyproject.toml

@@ -1,20 +1,26 @@
 [tool.poetry]
-name = "main"
+name = "testdata"
 version = "0.1.0"
 description = ""
-authors = ["Your Name <you@example.com>"]
+authors = ["Kristian Krsnik <git@krsnik.at>"]
 readme = "README.md"
-packages = [{ include = "src" }]
+packages = [{ include = "testdata" }]
 
 [tool.poetry.dependencies]
 python = "^3.11"
-fastapi = "^0.110.1"
-hypercorn = "^0.16.0"
 pydantic = "^2.6.4"
 ipaddress = "^1.0.23"
+uvicorn = "^0.30.3"
+python-fasthtml = "^0.2.1"
 
 [tool.poetry.scripts]
-testdata = "src.main:main"
+testdata = "testdata.main:main"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.2"
+requests = "^2.32.3"
+types-requests = "^2.32.0.20240712"
+pylint = "^3.2.6"
 
 [build-system]
 requires = ["poetry-core"]

src/main.py

@@ -1,114 +0,0 @@
-import sys
-import asyncio
-import argparse
-import json
-from os.path import exists
-
-from fastapi import FastAPI, Request, HTTPException, Query
-from fastapi.responses import StreamingResponse
-from fastapi import status
-from hypercorn.config import Config
-from hypercorn.asyncio import serve
-import ipaddress
-
-from .utils import convert_to_bytes, generate_data, load_database, save_database
-
-# Setup Parser
-parser = argparse.ArgumentParser()
-parser.add_argument('-c', '--config', type=argparse.FileType('r'),
-                    default='./config.json', help='Path to config file in JSON format.')
-
-args = parser.parse_args(sys.argv[1:])
-
-# Load Config
-CONFIG = json.load(args.config)
-BUFFER_SIZE = convert_to_bytes(CONFIG['buffer-size'])
-MAX_SIZE = convert_to_bytes(CONFIG['max-size'])
-MAX_DATA = convert_to_bytes(CONFIG['max-data'])
-AUTHORIZED_KEYS = CONFIG['keys']
-if isinstance(AUTHORIZED_KEYS, str):
-    with open(AUTHORIZED_KEYS) as file:
-        AUTHORIZED_KEYS = list(
-            filter(lambda x: x.strip() != '', file.read().splitlines()))
-DATABASE = CONFIG['database']
-
-if not exists(DATABASE):
-    save_database(DATABASE, {'data-used': 0})
-
-
-api = FastAPI(docs_url=None, redoc_url=None)
-
-
-class MaxSizePerRequestError(Exception):
-    pass
-
-
-class MinSizePerRequestError(Exception):
-    pass
-
-
-@api.get('/zeros')
-async def test_data(api_key: str, size: str) -> StreamingResponse:
-    try:
-        if api_key not in AUTHORIZED_KEYS:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail='Invalid API Key.'
-            )
-
-        try:
-            size = convert_to_bytes(size)
-        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
-        elif MAX_SIZE < size:
-            raise MaxSizePerRequestError
-
-        database = load_database(DATABASE)
-        if MAX_DATA <= database['data-used'] + size:
-            raise HTTPException(
-                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-                detail='Service not available.'
-            )
-        database['data-used'] += size
-
-        save_database(DATABASE, database)
-
-        return StreamingResponse(
-            status_code=status.HTTP_200_OK,
-            content=generate_data(size, BUFFER_SIZE),
-            media_type='application/octet-stream',
-            headers={
-                'Content-Length': str(size)
-            }
-        )
-
-    except MinSizePerRequestError as err:
-        raise HTTPException(
-            status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
-            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():
-    asyncio.run(serve(
-        api,
-        Config().from_mapping(
-            bind=CONFIG['binds'],
-            accesslog='-'
-        )
-    ))
-
-
-if __name__ == '__main__':
-    main()

test/test_imports.py

@@ -0,0 +1,3 @@
+def test_import():
+    import testdata
+    from testdata import run

test/test_run.py

@@ -0,0 +1,96 @@
+from multiprocessing import Process
+import os
+import json
+import tempfile
+import random
+
+import pytest
+import requests
+
+import testdata
+
+
+PROTOCOL = 'http'
+HOST = '127.0.0.1'
+PORT = 8080
+BUFFER_SIZE = 4 * 1024
+MAX_SIZE = 2 * 1024 * 1024 * 1024
+MAX_DATA = 10 * 1024 * 1024 * 1024
+API_KEY = f'{random.randrange(16 ** 32):032x}'
+API_KEYS = { API_KEY }
+
+TIMEOUT = 5
+
+@pytest.fixture(autouse = True)
+def server():
+    # Create Temporary Databases File
+    database = tempfile.NamedTemporaryFile(delete = False).name
+
+    proc = Process(target = testdata.run, args = (HOST, PORT, API_KEYS, MAX_SIZE, MAX_DATA, database, BUFFER_SIZE))
+
+    proc.start()
+
+    # Wait until webserver becomes available
+    while True:
+        try:
+            requests.get(f'{PROTOCOL}://{HOST}:{PORT}', timeout = TIMEOUT)
+        except requests.ConnectionError:
+            continue
+
+        break
+
+    yield database
+
+    # Terminate webserver
+    proc.terminate()
+    proc.join()
+
+    # Delete Temporary File
+    os.unlink(database)
+
+
+def test_get_file():
+
+    response = requests.get(
+        f'{PROTOCOL}://{HOST}:{PORT}/zeros/?api_key={API_KEY}&size=32', timeout=TIMEOUT)
+    assert response.content == b'\0' * 32
+
+
+def test_get_file_B():
+
+    response = requests.get(
+        f'{PROTOCOL}://{HOST}:{PORT}/zeros/?api_key={API_KEY}&size=32B', timeout=TIMEOUT)
+    assert response.content == b'\0' * 32
+
+
+def test_get_file_KB():
+
+    response = requests.get(
+        f'{PROTOCOL}://{HOST}:{PORT}/zeros/?api_key={API_KEY}&size=32KB', timeout=TIMEOUT)
+    assert response.status_code == 200
+    assert response.content == b'\0' * 32 * 1000
+
+
+def test_get_file_KiB():
+
+    response = requests.get(
+        f'{PROTOCOL}://{HOST}:{PORT}/zeros/?api_key={API_KEY}&size=32KiB', timeout=TIMEOUT)
+    assert response.status_code == 200
+    assert response.content == b'\0' * 32 * 1024
+
+
+def test_get_file_invalid_format():
+
+    response = requests.get(
+        f'{PROTOCOL}://{HOST}:{PORT}/zeros/?api_key={API_KEY}&size=32Invalid', timeout=TIMEOUT)
+    assert response.status_code == 400
+    assert response.text == 'Invalid Format.'
+
+
+def test_database_data_used(server):
+
+    requests.get(
+        f'{PROTOCOL}://{HOST}:{PORT}/zeros/?api_key={API_KEY}&size=32', timeout=TIMEOUT)
+    with open(server) as file:
+        assert json.loads(file.read())['data-used'] == 32
+

test/test_utils.py

@@ -0,0 +1,23 @@
+from pytest import raises
+
+from testdata.utils import convert_to_bytes
+
+
+def test_convert_to_bytes():
+    test_values = {
+        ('0', 0),
+        ('32', 32),
+        ('10_000', 10000),
+        ('999999999999999999', 999999999999999999),
+    }
+
+    for input, expected_output in test_values:
+        assert convert_to_bytes(input) == expected_output
+
+    test_exceptions = {
+        ('9E3', ValueError)
+    }
+
+    for input, exception in test_exceptions:
+        with raises(exception):
+            convert_to_bytes(input)

testdata/__init__.py

@@ -0,0 +1 @@
+from .main import run

testdata/__main__.py

@@ -0,0 +1,3 @@
+from testdata.main import main
+
+main()

testdata/app.py

@@ -0,0 +1,124 @@
+# pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
+# pylint: disable=line-too-long
+
+from pathlib import Path
+
+from fasthtml.fastapp import FastHTML, FastHTMLWithLiveReload, Request, StaticFiles, picolink
+from fasthtml.common import Form, Label, Input, Br, Button, A, Html, serve, Body, Link, P, Body, Div, Span, Main, Title # pylint: disable=no-name-in-module
+from starlette import status
+from starlette.responses import Response, StreamingResponse
+from starlette.exceptions import HTTPException
+from pydantic import ValidationError
+
+from .utils import load_database, save_database, generate_data
+from .custom_types import TestDataBody, Config
+
+
+class MaxSizePerRequestError(Exception):
+    pass
+
+
+class MinSizePerRequestError(Exception):
+    pass
+
+
+with open('./config.json') as file:  # parse_cli_arguments(sys.argv[1:])
+    config = Config.model_validate_json(file.read())
+
+
+app = FastHTMLWithLiveReload(hdrs=(
+    (Link(rel='stylesheet', type='text/css', href='static/style.css'),)
+))
+
+app.mount(
+    '/static',
+    StaticFiles(directory=Path(__file__).parent.absolute() / 'static'),
+    name='static'
+)
+
+
+@app.get('/')
+async def homepage():
+    return Main(
+        P('Testdata Download'),
+        P('Download a file of exact size. The file will only contain zeros.'),
+        Div(
+            Div(
+                Form(
+                    Div(
+                        Label('API Key'), Br(),
+                        Input(name='api_key'), Br()
+                    ),
+                    Div(
+                        Label('Size'), Br(),
+                        Input(name='size', placeholder='2MiB, 1GB, etc.'), Br()
+                    ),
+                    Button('Download'),
+                    action='/zeros',  hx_boost="true"
+                ),
+                P(f'Max Download Size: {config.max_size / (1024 * 1024 * 1024):.02f} GiB'), Br(),  # nopep8
+                cls='download'
+            ),
+            P('Contact: ', A('contact.testdata@krsnik.at', href='mailto:contact.testdata@krsnik.at'), id='contact')  # nopep8
+        )
+    )
+
+
+@ app.get('/zeros')  # type: ignore
+async def get_zeros(api_key: str, size: str, request: Request) -> StreamingResponse:
+    # AJAX Requests cannot return a file, thus send a redirect.
+    if 'HX-Request' in request.headers:
+        return Response(
+            content='<p>Hello</p>',
+            headers={'HX-Redirect': f'/zeros/?api_key={api_key}&size={size}'})
+
+    try:
+        body = TestDataBody(api_key=api_key, size=size)  # type: ignore
+    except ValidationError as err:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail='Invalid Format.'
+        ) from err
+
+    try:
+        if body.api_key not in config.api_keys:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail='Invalid API Key.'
+            )
+
+        if body.size < 0:
+            raise MinSizePerRequestError
+
+        if config.max_size < body.size:
+            raise MaxSizePerRequestError
+
+        db = load_database(config.database)
+        if config.max_data <= db['data-used'] + body.size:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail='Service not available.'
+            )
+
+        db['data-used'] += body.size
+        save_database(config.database, db)
+
+        return StreamingResponse(
+            status_code=status.HTTP_200_OK,
+            content=generate_data(body.size, config.buffer_size),
+            media_type='application/octet-stream',
+            headers={
+                'Content-Length': str(body.size),
+            }
+        )
+
+    except MinSizePerRequestError as err:
+        raise HTTPException(
+            status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
+            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 {config.max_size} Bytes.'
+        ) from err

testdata/custom_types.py

@@ -0,0 +1,41 @@
+from pydantic import BaseModel, Field, ConfigDict, field_validator
+
+from .utils import convert_to_bytes
+
+
+class Config(BaseModel):
+    host: str
+    port: int
+    buffer_size: int = Field(alias='buffer-size', default=4 * 1024)  # 4KB
+    max_size: int = Field(alias='max-size', default=2 * 1024 ** 3)  # 2GB
+    max_data: int = Field(alias='max-data', default=0)  # unlimited
+    api_keys: set[str] = Field(alias='api-keys')
+    database: str
+    log: str = '-'
+
+    model_config = ConfigDict(extra='forbid')
+
+    @field_validator('buffer_size', 'max_size', 'max_data', mode='before')
+    @classmethod
+    def convert_size(cls, value: int | str) -> int:
+        return convert_to_bytes(value)
+
+    @field_validator('api_keys', mode='before')
+    @classmethod
+    def check_if_key_in_file(cls, value: str | list[str]) -> set[str]:
+        if isinstance(value, str):
+            with open(value) as file:
+                return set(filter(lambda x: x.strip() != '', file.read().splitlines()))
+        return set(value)
+
+
+class TestDataBody(BaseModel):
+    api_key: str
+    size: int
+
+    model_config = ConfigDict(extra='forbid')
+
+    @field_validator('size', mode='before')
+    @classmethod
+    def convert_size(cls, value: str) -> int:
+        return convert_to_bytes(value)

testdata/main.py

@@ -0,0 +1,38 @@
+import sys
+import argparse
+import os
+from os.path import exists
+
+import uvicorn
+
+from .utils import save_database
+from .custom_types import Config
+
+
+# Setup Parser
+def parse_cli_arguments(argv: list[str]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '-c',
+        '--config',
+        type = argparse.FileType('r'),
+        default = './config.json',
+        help = 'Path to config file in JSON format.'
+    )
+
+    return parser.parse_args(argv)
+
+
+def run(host: str, port: int, api_keys: set[str], max_size: int, max_data: int, database: str, buffer_size: int, reload: bool = False):
+    if not exists(database) or os.stat(database).st_size == 0:
+        save_database(database, {'data-used': 0})
+
+    # app = create_api(api_keys, max_size, max_data, database, buffer_size)
+
+    uvicorn.run("testdata.app:app", host=host, port=port, reload=reload)
+
+
+def main():
+    args = parse_cli_arguments(sys.argv[1:])
+    config = Config.model_validate_json(args.config.read())
+    run(**config.model_dump(exclude={'log'}), reload=True)

testdata/static/style.css

@@ -0,0 +1,48 @@
+@font-face {
+    font-family: roboto;
+    src: url(Roboto-Regular.ttf);
+}
+
+:root {
+    --background-color: #ffffff;
+    --form-background-color: #e2ddff;
+}
+
+body {
+    background-color: var(--background-color);
+    font-family: roboto;
+}
+
+.download {
+    background-color: var(--form-background-color);
+    border-radius: 1em;
+    text-align: center;
+    width: 70%;
+    margin: 0 auto;
+    margin-top: 5em;
+}
+
+.download > form > div {
+    width: 70%;
+    margin: 0 auto;
+}
+
+.download > form > div > label {
+    font-size: 3em;
+}
+
+.download > form > div > input {
+    width: 100%;
+    font-size: 3em;
+    text-indent: 1em;
+    background-color: var(--background-color);
+}
+
+.download > form > button {
+    margin-top: 1em;
+    font-size: 3em;
+}
+
+#contact {
+    text-align: center;
+}

testdata/utils.py

@@ -1,10 +1,15 @@
 import json
 import asyncio
+from typing import AsyncGenerator
+
 
 def convert_to_bytes(size: int | str) -> int:
+    if isinstance(size, int):
+        return size
+
     try:
         return int(size)
-    except ValueError:  # treat as string
+    except ValueError as err:
         units = {
             'TB': 1000 ** 4, 'TiB': 1024 ** 4,
             'GB': 1000 ** 3, 'GiB': 1024 ** 3,
@@ -13,15 +18,16 @@ def convert_to_bytes(size: int | str) -> int:
             'B': 1
         }
 
-        for unit in units:
+        for unit, value in units.items():
             if size.endswith(unit):
-                return int(float(size.removesuffix(unit)) * units[unit])
-                break
+                return int(float(size.removesuffix(unit)) * value)
 
-        raise ValueError
+        raise ValueError(
+            'Invalid format. Expected integer or float ending with a data unit (B, KB, MiB,...).'
+        ) from err
 
 
-async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
+async def generate_data(size: int, buffer_size: int = 4 * 1024) -> AsyncGenerator[bytes, None]:
     size_left = size
 
     # https://github.com/tiangolo/fastapi/issues/5183
@@ -32,23 +38,18 @@ async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
             size_left -= buffer_size
             yield b'\0' * buffer_size
             await asyncio.sleep(0)
-        else:
+
         yield b'\0' * size_left
         await asyncio.sleep(0)
-    except asyncio.CancelledError:
-        raise GeneratorExit
-
-
-def check_policies(ip: str) -> None:
-    network = ipaddress.ip_network(ip)
-    print(network)
+    except asyncio.CancelledError as err:
+        raise GeneratorExit from err
 
 
 def load_database(path: str) -> dict:
-    with open(path, 'r') as file:
+    with open(path, 'r', encoding='utf-8') as file:
         return json.load(file)
 
 
 def save_database(path: str, database: dict) -> None:
-    with open(path, 'w') as file:
+    with open(path, 'w', encoding='utf-8') as file:
         json.dump(database, file, indent=2)