implemented some pydandic models and fixed type errors

This commit is contained in:
Kristian Krsnik 2024-08-03 00:52:27 +02:00
parent 93c51a00a0
commit 3f8323d297
Signed by: Kristian
GPG Key ID: FD1330AC9F909E85
8 changed files with 98 additions and 36 deletions

1
.gitignore vendored
View File

@ -8,3 +8,4 @@
# python
__pycache__/
.pytest_cache/
.mypy_cache/

View File

@ -23,7 +23,8 @@
packages = forAllSystems (system: {
default = poetry2nix.${system}.mkPoetryApplication {
projectDir = self;
checkPhase = "pytest";
checkPhase = "mypy . && pytest";
preferWheels = true; # avoid long build process for `mypy`
};
vm = self.nixosConfigurations.vm.config.microvm.declaredRunner;
});

16
poetry.lock generated
View File

@ -504,6 +504,20 @@ anyio = ">=3.4.0,<5"
[package.extras]
full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart (>=0.0.7)", "pyyaml"]
[[package]]
name = "types-requests"
version = "2.32.0.20240712"
description = "Typing stubs for requests"
optional = false
python-versions = ">=3.8"
files = [
{file = "types-requests-2.32.0.20240712.tar.gz", hash = "sha256:90c079ff05e549f6bf50e02e910210b98b8ff1ebdd18e19c873cd237737c1358"},
{file = "types_requests-2.32.0.20240712-py3-none-any.whl", hash = "sha256:f754283e152c752e46e70942fa2a146b5bc70393522257bb85bd1ef7e019dcc3"},
]
[package.dependencies]
urllib3 = ">=2"
[[package]]
name = "typing-extensions"
version = "4.12.2"
@ -553,4 +567,4 @@ standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)",
[metadata]
lock-version = "2.0"
python-versions = "^3.11"
content-hash = "a1f2445b7578ac1606a1968f771da2d0c31e859298a4d8815cfbfcf94d06ac57"
content-hash = "8c95de1de17ced501a0460661e22c77839a3afca88512e07ad95344e719679a5"

View File

@ -23,6 +23,7 @@ testdata = "testdata.main:main"
pytest = "^8.3.2"
requests = "^2.32.3"
mypy = "^1.11.1"
types-requests = "^2.32.0.20240712"
[build-system]
requires = ["poetry-core"]

38
testdata/api.py vendored
View File

@ -1,14 +1,14 @@
from fastapi import FastAPI, HTTPException, status
from fastapi.responses import StreamingResponse
from pydantic import ValidationError
from .utils import convert_to_bytes, load_database, save_database, generate_data
from .utils import load_database, save_database, generate_data
from .custom_types import TestDataBody
def create_api(api_keys: {str}, max_size: int, max_data: int, database: str, buffer_size: int):
def create_api(api_keys: set[str], max_size: int, max_data: int, database: str, buffer_size: int):
api = FastAPI(docs_url=None, redoc_url=None)
class MaxSizePerRequestError(Exception):
pass
@ -20,41 +20,41 @@ def create_api(api_keys: {str}, max_size: int, max_data: int, database: str, buf
@api.get('/')
async def test_data(api_key: str, size: str) -> StreamingResponse:
try:
if api_key not in api_keys:
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 api_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:
if body.size < 0:
raise MinSizePerRequestError
elif max_size < size:
elif max_size < body.size:
raise MaxSizePerRequestError
db = load_database(database)
if max_data <= db['data-used'] + size:
if max_data <= db['data-used'] + body.size:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail='Service not available.'
)
db['data-used'] += size
db['data-used'] += body.size
save_database(database, db)
return StreamingResponse(
status_code=status.HTTP_200_OK,
content=generate_data(size, buffer_size),
content=generate_data(body.size, buffer_size),
media_type='application/octet-stream',
headers={
'Content-Length': str(size)
'Content-Length': str(body.size)
}
)

36
testdata/custom_types.py vendored Normal file
View File

@ -0,0 +1,36 @@
from pydantic import BaseModel, ConfigDict, field_validator
from .utils import convert_to_bytes
# class Config(BaseModel):
# host: str
# port: int
# buffer_size: int = 4 * 1024 # 4KB
# max_size: int = 2 * 1024 * 1024 * 1024 # 2GB
# max_data: int = 0 # unlimited
# api_keys = set[str]
# database = str
# model_config = ConfigDict(extra='allow')
# @field_validator('buffer_size', 'max_size', 'max_data')
# @classmethod
# def convert_size(cls, value: int | str) -> int:
# return convert_to_bytes(value)
# TODO
# class DataBase(BaseModel):
# model_config = ConfigDict(extra='forbid')
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)

20
testdata/main.py vendored
View File

@ -9,13 +9,22 @@ import uvicorn
from .utils import convert_to_bytes, save_database
from .api import create_api
# 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.')
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: {str}, max_size: int, max_data: int, database: str, buffer_size: int):
def run(host: str, port: int, api_keys: set[str], max_size: int, max_data: int, database: str, buffer_size: int):
if not exists(database) or os.stat(database).st_size == 0:
save_database(database, {'data-used': 0})
@ -25,8 +34,9 @@ def run(host: str, port: int, api_keys: {str}, max_size: int, max_data: int, dat
def main():
args = parser.parse_args(sys.argv[1:])
args = parse_cli_arguments(sys.argv[1:])
config = json.load(args.config)
host = config['host']
port = config['port']
buffer_size = convert_to_bytes(config['buffer-size'])

19
testdata/utils.py vendored
View File

@ -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:
units = {
'TB': 1000 ** 4, 'TiB': 1024 ** 4,
'GB': 1000 ** 3, 'GiB': 1024 ** 3,
@ -16,12 +21,11 @@ def convert_to_bytes(size: int | str) -> int:
for unit in units:
if size.endswith(unit):
return int(float(size.removesuffix(unit)) * units[unit])
break
raise ValueError
else:
raise ValueError('Invalid format. Expected integer or float ending with a data unit (B, KB, MiB,...).')
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
@ -39,11 +43,6 @@ async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
raise GeneratorExit
def check_policies(ip: str) -> None:
network = ipaddress.ip_network(ip)
print(network)
def load_database(path: str) -> dict:
with open(path, 'r') as file:
return json.load(file)