Compare commits
1 Commits
Author | SHA1 | Date | |
---|---|---|---|
d9d8759600
|
18
.gitignore
vendored
18
.gitignore
vendored
@ -1,9 +1,3 @@
|
||||
# Nix #
|
||||
/result
|
||||
|
||||
# Direnv #
|
||||
/.direnv/
|
||||
|
||||
# Python #
|
||||
# Virtual Environment
|
||||
/.venv/
|
||||
@ -19,6 +13,18 @@ __pycache__/
|
||||
/.pytest_cache/
|
||||
/.mypy_cache/
|
||||
|
||||
# Nix #
|
||||
|
||||
# Build
|
||||
/result
|
||||
|
||||
# MicroVM
|
||||
/var.img
|
||||
/control.socket
|
||||
|
||||
# Direnv #
|
||||
/.direnv/
|
||||
|
||||
# Project specific files #
|
||||
config.json
|
||||
db.json
|
||||
|
28
Dockerfile
28
Dockerfile
@ -1,28 +0,0 @@
|
||||
# Nix builder
|
||||
FROM nixos/nix:latest AS builder
|
||||
|
||||
# Copy our source and setup our working dir.
|
||||
COPY . /tmp/build
|
||||
WORKDIR /tmp/build
|
||||
|
||||
# Build our Nix environment
|
||||
RUN nix \
|
||||
--extra-experimental-features "nix-command flakes" \
|
||||
--option filter-syscalls false \
|
||||
build
|
||||
|
||||
# Copy the Nix store closure into a directory. The Nix store closure is the
|
||||
# entire set of Nix store values that we need for our build.
|
||||
RUN mkdir /tmp/nix-store-closure
|
||||
RUN cp -r $(nix-store -qR result/) /tmp/nix-store-closure
|
||||
|
||||
# Final image is based on scratch. We copy a bunch of Nix dependencies
|
||||
# but they're fully self-contained so we don't need Nix anymore.
|
||||
FROM scratch
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Copy /nix/store
|
||||
COPY --from=builder /tmp/nix-store-closure /nix/store
|
||||
COPY --from=builder /tmp/build/result /app
|
||||
CMD ["/app/bin/testdata"]
|
20
README.md
20
README.md
@ -1,15 +1,19 @@
|
||||
# Simple Testdata Generator
|
||||
# Simple Test Data Generator
|
||||
|
||||
## Example Config
|
||||
|
||||
```json
|
||||
{
|
||||
"keys": ["TESTKEY1", "TESTKEY2", "TESTKEY3"],
|
||||
"max-size": "1GiB",
|
||||
"max-data": "1TiB",
|
||||
"buffer-size": "12MiB",
|
||||
"database": "./db.json",
|
||||
"database-update-interval": 5.0,
|
||||
"log": "./log.jsonl"
|
||||
"binds": [
|
||||
"127.0.0.1:9250"
|
||||
],
|
||||
"log": "-",
|
||||
"buffer-size": "4KiB",
|
||||
"max-size": "2GB",
|
||||
"keys": [
|
||||
"TESTKEY"
|
||||
],
|
||||
"max-data": "10GB",
|
||||
"database": "database.json"
|
||||
}
|
||||
```
|
||||
|
@ -1,15 +0,0 @@
|
||||
services:
|
||||
testdata:
|
||||
image: result/latest
|
||||
|
||||
build:
|
||||
dockerfile: ./Dockerfile
|
||||
|
||||
environment:
|
||||
TESTDATA_HOST: 0.0.0.0
|
||||
TESTDATA_PORT: 1234
|
||||
TESTDATA_CONFIG: ./config.json
|
||||
volumes:
|
||||
- ./config.json:/app/config.json
|
||||
- ./db.json:/app/db.json
|
||||
- ./log.jsonl:/app/log.jsonl
|
24
flake.lock
generated
24
flake.lock
generated
@ -27,11 +27,11 @@
|
||||
"spectrum": "spectrum"
|
||||
},
|
||||
"locked": {
|
||||
"lastModified": 1748464257,
|
||||
"narHash": "sha256-PdnQSE2vPfql9WEjunj2qQnDpuuvk7HH+4djgXJSwFs=",
|
||||
"lastModified": 1735074045,
|
||||
"narHash": "sha256-CeYsC8J2dNiV2FCQOxK1oZ/jNpOF2io7aCEFHmfi95U=",
|
||||
"owner": "astro",
|
||||
"repo": "microvm.nix",
|
||||
"rev": "e238645b6f0447a2eb1d538d300d5049d4006f9f",
|
||||
"rev": "2ae08de8e8068b00193b9cfbc0acc9dfdda03181",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
@ -42,15 +42,15 @@
|
||||
},
|
||||
"nixpkgs": {
|
||||
"locked": {
|
||||
"lastModified": 1748370509,
|
||||
"narHash": "sha256-QlL8slIgc16W5UaI3w7xHQEP+Qmv/6vSNTpoZrrSlbk=",
|
||||
"owner": "NixOS",
|
||||
"lastModified": 1735471104,
|
||||
"narHash": "sha256-0q9NGQySwDQc7RhAV2ukfnu7Gxa5/ybJ2ANT8DQrQrs=",
|
||||
"owner": "nixos",
|
||||
"repo": "nixpkgs",
|
||||
"rev": "4faa5f5321320e49a78ae7848582f684d64783e9",
|
||||
"rev": "88195a94f390381c6afcdaa933c2f6ff93959cb4",
|
||||
"type": "github"
|
||||
},
|
||||
"original": {
|
||||
"owner": "NixOS",
|
||||
"owner": "nixos",
|
||||
"ref": "nixos-unstable",
|
||||
"repo": "nixpkgs",
|
||||
"type": "github"
|
||||
@ -65,11 +65,11 @@
|
||||
"spectrum": {
|
||||
"flake": false,
|
||||
"locked": {
|
||||
"lastModified": 1746869549,
|
||||
"narHash": "sha256-BKZ/yZO/qeLKh9YqVkKB6wJiDQJAZNN5rk5NsMImsWs=",
|
||||
"lastModified": 1733308308,
|
||||
"narHash": "sha256-+RcbMAjSxV1wW5UpS9abIG1lFZC8bITPiFIKNnE7RLs=",
|
||||
"ref": "refs/heads/main",
|
||||
"rev": "d927e78530892ec8ed389e8fae5f38abee00ad87",
|
||||
"revCount": 862,
|
||||
"rev": "80c9e9830d460c944c8f730065f18bb733bc7ee2",
|
||||
"revCount": 792,
|
||||
"type": "git",
|
||||
"url": "https://spectrum-os.org/git/spectrum"
|
||||
},
|
||||
|
134
flake.nix
134
flake.nix
@ -2,7 +2,7 @@
|
||||
description = "A webserver to create files for testing purposes";
|
||||
|
||||
inputs = {
|
||||
nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
|
||||
nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
|
||||
|
||||
microvm = {
|
||||
url = "github:astro/microvm.nix";
|
||||
@ -10,25 +10,82 @@
|
||||
};
|
||||
};
|
||||
|
||||
outputs =
|
||||
{
|
||||
outputs = {
|
||||
self,
|
||||
nixpkgs,
|
||||
...
|
||||
}@inputs:
|
||||
let
|
||||
supportedSystems = [ "x86_64-linux" ];
|
||||
} @ inputs: let
|
||||
supportedSystems = ["x86_64-linux" "x86_64-darwin" "aarch64-linux" "aarch64-darwin"];
|
||||
forAllSystems = nixpkgs.lib.genAttrs supportedSystems;
|
||||
pkgs = forAllSystems (system: nixpkgs.legacyPackages.${system});
|
||||
pkgs = forAllSystems (system: nixpkgs.legacyPackages.${system}.extend overlay);
|
||||
|
||||
overlay = final: prev: rec {
|
||||
python3Packages = prev.python3Packages.overrideScope (pfinal: pprev: {
|
||||
packageNameToDrv = x: builtins.getAttr (cleanPythonPackageName x) final.python3Packages;
|
||||
});
|
||||
|
||||
cleanPythonPackageName = x: let
|
||||
cleanName = builtins.match "([a-z,A-Z,0-9,_,-]+).*" x;
|
||||
in
|
||||
{
|
||||
if cleanName != null
|
||||
then builtins.elemAt cleanName 0
|
||||
else builtins.warn "Could not determine package name from '${x}'" null;
|
||||
};
|
||||
|
||||
pyproject = builtins.fromTOML (builtins.readFile ./pyproject.toml);
|
||||
|
||||
buildDependencies = forAllSystems (system: builtins.map pkgs.${system}.python3Packages.packageNameToDrv pyproject.build-system.requires);
|
||||
runtimeDependencies = forAllSystems (system: builtins.map pkgs.${system}.python3Packages.packageNameToDrv pyproject.project.dependencies);
|
||||
optionalDependencies = forAllSystems (system: builtins.mapAttrs (name: value: builtins.map pkgs.${system}.python3Packages.packageNameToDrv value) pyproject.project.optional-dependencies);
|
||||
in {
|
||||
# `nix build`
|
||||
packages = forAllSystems (system: rec {
|
||||
packages = forAllSystems (system: let
|
||||
buildTestdata = {skipCheck ? false}:
|
||||
pkgs.${system}.python3Packages.buildPythonPackage {
|
||||
pname = pyproject.project.name;
|
||||
version = pyproject.project.version;
|
||||
src = ./.;
|
||||
|
||||
pyproject = true;
|
||||
|
||||
build-system = buildDependencies.${system};
|
||||
|
||||
dependencies = runtimeDependencies.${system};
|
||||
|
||||
optional-dependencies = optionalDependencies.${system};
|
||||
|
||||
nativeCheckInputs = optionalDependencies.${system}.dev;
|
||||
|
||||
checkPhase = let
|
||||
dev = builtins.map (x: x.pname) optionalDependencies.${system}.dev;
|
||||
in ''
|
||||
${
|
||||
if builtins.elem "pytest" dev && !skipCheck
|
||||
then "pytest src tests"
|
||||
else ""
|
||||
}
|
||||
${
|
||||
if builtins.elem "mypy" dev && !skipCheck
|
||||
then "mypy src tests"
|
||||
else ""
|
||||
}
|
||||
${
|
||||
if builtins.elem "pylint" dev && !skipCheck
|
||||
then "pylint src tests"
|
||||
else ""
|
||||
}
|
||||
'';
|
||||
};
|
||||
in rec {
|
||||
default = testdata;
|
||||
testdata = pkgs.${system}.callPackage ./nix/package.nix { src = ./.; };
|
||||
testdata = buildTestdata {skipCheck = false;};
|
||||
quick = buildTestdata {skipCheck = true;};
|
||||
vm = self.nixosConfigurations.vm.config.microvm.declaredRunner;
|
||||
});
|
||||
|
||||
# `nix fmt`
|
||||
formatter = forAllSystems (system: pkgs.${system}.alejandra);
|
||||
|
||||
# `nix develop`
|
||||
devShells = forAllSystems (system: rec {
|
||||
default = venv;
|
||||
@ -51,24 +108,36 @@
|
||||
# NixOS Module
|
||||
nixosModules.default = import ./nix/module.nix inputs;
|
||||
|
||||
# NixOS definition for a microvm to test nixosModules
|
||||
nixosConfigurations."vm" = nixpkgs.lib.nixosSystem {
|
||||
# nixos definition for a microvm to test nixosModules
|
||||
nixosConfigurations = let
|
||||
system = "x86_64-linux";
|
||||
in {
|
||||
vm = nixpkgs.lib.nixosSystem {
|
||||
inherit system;
|
||||
|
||||
modules = [
|
||||
inputs.microvm.nixosModules.microvm
|
||||
(
|
||||
{ config, ... }:
|
||||
{
|
||||
services.getty.autologinUser = "root";
|
||||
({config, ...}: {
|
||||
system.stateVersion = config.system.nixos.version;
|
||||
|
||||
networking.hostName = "vm";
|
||||
users.users.root.password = "";
|
||||
|
||||
microvm = {
|
||||
hypervisor = "qemu";
|
||||
|
||||
# volumes = [
|
||||
# {
|
||||
# mountPoint = "/var";
|
||||
# image = "var.img";
|
||||
# size = 256;
|
||||
# }
|
||||
# ];
|
||||
shares = [
|
||||
{
|
||||
# Host's /nix/store will be picked up so that no squashfs/erofs will be built for it.
|
||||
# use proto = "virtiofs" for MicroVMs that are started by systemd
|
||||
proto = "9p";
|
||||
tag = "ro-store";
|
||||
# a host's /nix/store will be picked up so that no
|
||||
# squashfs/erofs will be built for it.
|
||||
source = "/nix/store";
|
||||
mountPoint = "/nix/.ro-store";
|
||||
}
|
||||
@ -78,7 +147,6 @@
|
||||
{
|
||||
type = "user";
|
||||
id = "qemu";
|
||||
# Locally administered have one of 2/6/A/E in the second nibble.
|
||||
mac = "02:00:00:01:01:01";
|
||||
}
|
||||
];
|
||||
@ -89,13 +157,16 @@
|
||||
guest.port = config.services.testdata.port;
|
||||
}
|
||||
];
|
||||
|
||||
# "qemu" has 9p built-in!
|
||||
hypervisor = "qemu";
|
||||
socket = "control.socket";
|
||||
};
|
||||
}
|
||||
)
|
||||
})
|
||||
|
||||
self.nixosModules.default
|
||||
rec {
|
||||
networking.firewall.allowedTCPPorts = [ services.testdata.port ];
|
||||
networking.firewall.allowedTCPPorts = [services.testdata.port];
|
||||
services.testdata = {
|
||||
enable = true;
|
||||
|
||||
@ -103,21 +174,16 @@
|
||||
port = 1234;
|
||||
|
||||
settings = {
|
||||
keys = [
|
||||
"one"
|
||||
"two"
|
||||
"three"
|
||||
];
|
||||
max-size = "1GB";
|
||||
max-data = "100GB";
|
||||
buffer-size = "12MiB";
|
||||
database = "/root/testdata-state.json";
|
||||
database-update-interval = 5.0;
|
||||
log = "/root/log.jsonl";
|
||||
"keys" = ["one" "two" "three"];
|
||||
"max-size" = "1GB";
|
||||
"max-data" = "100GB";
|
||||
"buffer-size" = "12MiB";
|
||||
"database" = "/root/testdata_state.json";
|
||||
};
|
||||
};
|
||||
}
|
||||
];
|
||||
};
|
||||
};
|
||||
};
|
||||
}
|
||||
|
@ -1,32 +1,23 @@
|
||||
inputs:
|
||||
{
|
||||
inputs: {
|
||||
config,
|
||||
lib,
|
||||
pkgs,
|
||||
...
|
||||
}:
|
||||
let
|
||||
}: let
|
||||
cfg = config.services.testdata;
|
||||
package = inputs.self.packages.${pkgs.stdenv.hostPlatform.system}.default;
|
||||
inherit (lib)
|
||||
mkIf
|
||||
mkEnableOption
|
||||
mkOption
|
||||
types
|
||||
;
|
||||
inherit (lib) mkIf mkEnableOption mkOption types;
|
||||
|
||||
format = pkgs.formats.json { };
|
||||
format = pkgs.formats.json {};
|
||||
configFile = format.generate "config.json" cfg.settings;
|
||||
in
|
||||
{
|
||||
in {
|
||||
options.services.testdata = {
|
||||
enable = mkEnableOption "testdata";
|
||||
|
||||
settings = mkOption {
|
||||
type =
|
||||
with types;
|
||||
let
|
||||
type = with types; let
|
||||
valueType = nullOr (oneOf [
|
||||
# TODO: restrict type to actual config file structure
|
||||
bool
|
||||
int
|
||||
float
|
||||
@ -52,7 +43,7 @@ in
|
||||
};
|
||||
|
||||
config = mkIf cfg.enable {
|
||||
environment.systemPackages = [ package ];
|
||||
environment.systemPackages = [package];
|
||||
|
||||
systemd.services.testdata = {
|
||||
enable = true;
|
||||
@ -62,7 +53,7 @@ in
|
||||
ExecStart = "${package}/bin/testdata --config ${configFile} --listen ${cfg.host} --port ${builtins.toString cfg.port}";
|
||||
};
|
||||
|
||||
wantedBy = [ "multi-user.target" ];
|
||||
wantedBy = ["multi-user.target"];
|
||||
};
|
||||
};
|
||||
}
|
||||
|
@ -1,46 +0,0 @@
|
||||
{
|
||||
src,
|
||||
python3Packages,
|
||||
}:
|
||||
let
|
||||
inherit (python3Packages)
|
||||
setuptools
|
||||
fastapi
|
||||
uvicorn
|
||||
pydantic
|
||||
pytest
|
||||
requests
|
||||
mypy
|
||||
pylint
|
||||
;
|
||||
|
||||
project = (builtins.fromTOML (builtins.readFile "${src}/pyproject.toml")).project;
|
||||
pname = project.name;
|
||||
version = project.version;
|
||||
in
|
||||
python3Packages.buildPythonPackage {
|
||||
inherit pname version src;
|
||||
|
||||
pyproject = true;
|
||||
|
||||
build-system = [ setuptools ];
|
||||
|
||||
dependencies = [
|
||||
fastapi
|
||||
uvicorn
|
||||
pydantic
|
||||
];
|
||||
|
||||
nativeCheckInputs = [
|
||||
pytest
|
||||
requests
|
||||
mypy
|
||||
pylint
|
||||
];
|
||||
|
||||
checkPhase = ''
|
||||
pytest tests
|
||||
mypy src
|
||||
pylint src
|
||||
'';
|
||||
}
|
@ -1,23 +1,27 @@
|
||||
[project]
|
||||
name = "testdata"
|
||||
version = "1.2.2"
|
||||
requires-python = "~=3.12, <4"
|
||||
dependencies = ["fastapi~=0.115", "uvicorn~=0.32", "pydantic~=2.9"]
|
||||
version = "1.1.0"
|
||||
requires-python = ">=3.11, <4"
|
||||
dependencies = [
|
||||
"fastapi==0.115.*",
|
||||
"uvicorn==0.32.*",
|
||||
"pydantic==2.10.*",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest~=8.3",
|
||||
"mypy~=1.13",
|
||||
"pylint~=3.3",
|
||||
"requests~=2.32",
|
||||
"types-requests~=2.32",
|
||||
"pytest==8.3.*",
|
||||
"mypy==1.13.*",
|
||||
"pylint==3.3.*",
|
||||
"requests==2.32.*",
|
||||
"types-requests==2.32.*"
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
testdata = "testdata.main:main"
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools~=78.1"]
|
||||
requires = ["setuptools==75.*"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
@ -36,5 +40,5 @@ disable = [
|
||||
"missing-class-docstring",
|
||||
"missing-function-docstring",
|
||||
"too-few-public-methods",
|
||||
"broad-exception-caught",
|
||||
"broad-exception-caught"
|
||||
]
|
||||
|
2
src/testdata/logger/logger.py
vendored
2
src/testdata/logger/logger.py
vendored
@ -132,7 +132,7 @@ def generate_log_config(log_path: str | None = None) -> dict:
|
||||
'class': logging.handlers.RotatingFileHandler,
|
||||
'level': 'DEBUG',
|
||||
'formatter': 'json',
|
||||
'filename': log_path,
|
||||
'filename': 'log.jsonl',
|
||||
'maxBytes': 1024 * 1024 * 10, # 10 MiB
|
||||
'backupCount': 3
|
||||
}} if log_path is not None else {}),
|
||||
|
26
src/testdata/main.py
vendored
26
src/testdata/main.py
vendored
@ -1,32 +1,14 @@
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import asyncio
|
||||
import shutil
|
||||
|
||||
from .testdata import Testdata
|
||||
|
||||
def parse_args(args: list[str]):
|
||||
def formatter(prog):
|
||||
return argparse.ArgumentDefaultsHelpFormatter(prog, max_help_position=shutil.get_terminal_size().columns)
|
||||
|
||||
parser = argparse.ArgumentParser(formatter_class=formatter)
|
||||
|
||||
parser.add_argument(
|
||||
'-c', '--config', type=argparse.FileType('r'),
|
||||
default=os.environ['TESTDATA_CONFIG'] if 'TESTDATA_CONFIG' in os.environ else './config.json',
|
||||
help='Path to config file in JSON format.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-l', '--listen', type=str,
|
||||
default=os.environ['TESTDATA_HOST'] if 'TESTDATA_HOST' in os.environ else '0.0.0.0',
|
||||
help='IP on which to listen.'
|
||||
)
|
||||
parser.add_argument(
|
||||
'-p', '--port', type=int,
|
||||
default=os.environ['TESTDATA_PORT'] if 'TESTDATA_PORT' in os.environ else 8080,
|
||||
help='Port on which to serve the webserver.'
|
||||
)
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-c', '--config', type=argparse.FileType('r'), default='./config.json', help='Path to config file in JSON format.')
|
||||
parser.add_argument('-l', '--listen', type=str, default='0.0.0.0', help='IP on which to listen.')
|
||||
parser.add_argument('-p', '--port', type=int, default='8080', help='Port on which to serve the webserver.')
|
||||
|
||||
return parser.parse_args(args)
|
||||
|
||||
|
161
src/testdata/testdata.py
vendored
161
src/testdata/testdata.py
vendored
@ -1,16 +1,10 @@
|
||||
import os
|
||||
import json
|
||||
import asyncio
|
||||
import inspect
|
||||
import functools
|
||||
import random
|
||||
import importlib.metadata
|
||||
from datetime import datetime
|
||||
|
||||
import uvicorn
|
||||
from typing_extensions import Annotated
|
||||
from fastapi import FastAPI, Request, Security, status, HTTPException
|
||||
from fastapi.security import APIKeyHeader, APIKeyQuery
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, status, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel, ConfigDict, Field, BeforeValidator, ValidationError
|
||||
|
||||
@ -49,113 +43,37 @@ class Testdata:
|
||||
max_size: Annotated[int, BeforeValidator(to_bytes)] = Field(alias='max-size')
|
||||
max_data: Annotated[int, BeforeValidator(to_bytes)] = Field(alias='max-data')
|
||||
buffer_size: Annotated[int, BeforeValidator(to_bytes)] = Field(alias='buffer-size')
|
||||
database: str | None = None
|
||||
log: str | None = Field(alias='log', default=None)
|
||||
database_update_interval: float = Field(alias='database-update-interval', default=5)
|
||||
database: str
|
||||
log_path: str | None = Field(alias='log-path', default=None)
|
||||
update_database_interval: float = Field(alias='update-database-interval', default=5)
|
||||
|
||||
_config: Config
|
||||
_api: FastAPI
|
||||
_state: dict
|
||||
_state: dict[str, int]
|
||||
_logger: logger.Logger
|
||||
|
||||
def __init__(self, config: Config):
|
||||
self._config = config
|
||||
self._api = FastAPI(docs_url=None, redoc_url=None)
|
||||
self._logger = logger.getLogger('testdata')
|
||||
self._api = self._setup_api()
|
||||
|
||||
# Store internal state
|
||||
self._state = {
|
||||
'version': importlib.metadata.version('testdata'), # For future compatibility
|
||||
'data-used': {f'{(today := datetime.today()).year}-{today.month:02}': 0} # math each months data usage
|
||||
}
|
||||
self._state = {'data-used': 0}
|
||||
|
||||
def _setup_api(self) -> FastAPI:
|
||||
api = FastAPI(docs_url='/', redoc_url=None)
|
||||
|
||||
# Security
|
||||
def get_api_key(
|
||||
api_key_query: str = Security(APIKeyQuery(name="api_key", auto_error=False)),
|
||||
api_key_header: str = Security(APIKeyHeader(name="x-api-key", auto_error=False))
|
||||
) -> str:
|
||||
# https://joshdimella.com/blog/adding-api-key-auth-to-fast-api
|
||||
|
||||
if api_key_query in self._config.authorized_keys:
|
||||
return api_key_query
|
||||
if api_key_header in self._config.authorized_keys:
|
||||
return api_key_header
|
||||
@self._api.get('/zeros')
|
||||
async def zeros(api_key: str, size: int | str) -> StreamingResponse:
|
||||
try:
|
||||
self._logger.debug('', extra={'api_key': api_key, 'size': size})
|
||||
|
||||
if api_key not in config.authorized_keys:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail='Invalid or missing API Key'
|
||||
detail='Invalid API Key.'
|
||||
)
|
||||
|
||||
# A wrapper to set the function signature to accept the api key dependency
|
||||
def secure(func):
|
||||
# Get old signature
|
||||
positional_only, positional_or_keyword, variadic_positional, keyword_only, variadic_keyword = [], [], [], [], []
|
||||
for value in inspect.signature(func).parameters.values():
|
||||
if value.kind == inspect.Parameter.POSITIONAL_ONLY:
|
||||
positional_only.append(value)
|
||||
elif value.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
|
||||
positional_or_keyword.append(value)
|
||||
elif value.kind == inspect.Parameter.VAR_POSITIONAL:
|
||||
variadic_positional.append(value)
|
||||
elif value.kind == inspect.Parameter.KEYWORD_ONLY:
|
||||
keyword_only.append(value)
|
||||
elif value.kind == inspect.Parameter.VAR_KEYWORD:
|
||||
variadic_keyword.append(value)
|
||||
|
||||
# Avoid passing an unrecognized keyword
|
||||
if inspect.iscoroutinefunction(func):
|
||||
async def wrapper(*args, **kwargs):
|
||||
if len(variadic_keyword) == 0:
|
||||
if 'api_key' in kwargs:
|
||||
del kwargs['api_key']
|
||||
|
||||
return await func(*args, **kwargs)
|
||||
else:
|
||||
def wrapper(*args, **kwargs):
|
||||
if len(variadic_keyword) == 0:
|
||||
if 'api_key' in kwargs:
|
||||
del kwargs['api_key']
|
||||
|
||||
return func(*args, **kwargs)
|
||||
|
||||
# Override signature
|
||||
wrapper.__signature__ = inspect.signature(func).replace(
|
||||
parameters=(
|
||||
*positional_only,
|
||||
*positional_or_keyword,
|
||||
*variadic_positional,
|
||||
*keyword_only,
|
||||
inspect.Parameter('api_key', inspect.Parameter.POSITIONAL_OR_KEYWORD, default=Security(get_api_key)),
|
||||
*variadic_keyword
|
||||
)
|
||||
)
|
||||
|
||||
return functools.wraps(func)(wrapper)
|
||||
|
||||
# Routes
|
||||
api.get('/zeros')(secure(self._zeros))
|
||||
|
||||
return api
|
||||
|
||||
async def _zeros(self, size: int | str, request: Request, filename: str = 'zeros.bin') -> StreamingResponse:
|
||||
try:
|
||||
extra = {'id': f'{random.randint(0, 2 ** 32 - 1):08X}'}
|
||||
self._logger.debug(
|
||||
'Initiated request.',
|
||||
extra=extra | {
|
||||
'ip': request.client.host if request.client is not None else None,
|
||||
'query-params': dict(request.query_params),
|
||||
'headers': dict(request.headers)
|
||||
}
|
||||
)
|
||||
|
||||
try:
|
||||
size = convert_to_bytes(size)
|
||||
except ValueError as err:
|
||||
self._logger.warning('Invalid format for size.', extra=extra)
|
||||
self._logger.warning('Invalid format for size.', extra={'api_key': api_key, 'size': size})
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail='Invalid format for size.'
|
||||
@ -163,51 +81,47 @@ class Testdata:
|
||||
|
||||
if size < 0:
|
||||
raise MinSizePerRequestError
|
||||
if self._config.max_size < size:
|
||||
if config.max_size < size:
|
||||
raise MaxSizePerRequestError
|
||||
|
||||
# update internal state
|
||||
current_date = f'{(today := datetime.today()).year}-{today.month:02}'
|
||||
if current_date not in self._state['data-used']:
|
||||
self._state['data-used'][current_date] = 0
|
||||
if self._config.max_data < self._state['data-used'][current_date] + size:
|
||||
if config.max_data < self._state['data-used'] + size:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail='Service not available.'
|
||||
)
|
||||
self._state['data-used'][current_date] += size
|
||||
self._state['data-used'] += size
|
||||
|
||||
self._logger.debug('Successfully processed request.', extra=extra)
|
||||
return StreamingResponse(
|
||||
status_code=status.HTTP_200_OK,
|
||||
content=generate_data(size, self._config.buffer_size),
|
||||
content=generate_data(size, config.buffer_size),
|
||||
media_type='application/octet-stream',
|
||||
headers={
|
||||
'Content-Length': str(size),
|
||||
'Content-Disposition': f'attachment; filename="{filename}"'
|
||||
'Content-Length': str(size)
|
||||
}
|
||||
)
|
||||
|
||||
except MinSizePerRequestError as err:
|
||||
self._logger.warning('Size if negative.', extra=extra)
|
||||
self._logger.warning('Size if negative.', extra={'api_key': api_key, 'size': size})
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
|
||||
detail='Size has to be non-negative.'
|
||||
) from err
|
||||
except MaxSizePerRequestError as err:
|
||||
self._logger.warning('Exceeded max size per request.', extra=extra)
|
||||
self._logger.warning('Exceeded max size per request.', extra={'api_key': api_key, 'size': size})
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE,
|
||||
detail=f'Exceeded max size per request of {self._config.max_size} Bytes.'
|
||||
detail=f'Exceeded max size per request of {config.max_size} Bytes.'
|
||||
) from err
|
||||
except Exception as err:
|
||||
self._logger.exception(err)
|
||||
raise err
|
||||
|
||||
async def _update_state(self) -> None:
|
||||
assert self._config.database is not None
|
||||
|
||||
mode = 'r+' if os.path.exists(self._config.database) else 'w+'
|
||||
async def _update_state(self):
|
||||
if os.path.exists(self._config.database):
|
||||
mode = 'r+'
|
||||
else:
|
||||
mode = 'w+'
|
||||
|
||||
with open(self._config.database, mode, encoding='utf-8') as file:
|
||||
try:
|
||||
@ -217,23 +131,22 @@ class Testdata:
|
||||
|
||||
while True:
|
||||
file.seek(0)
|
||||
json.dump(self._state, file, indent=2)
|
||||
json.dump(self._state, file)
|
||||
file.truncate()
|
||||
await asyncio.sleep(self._config.database_update_interval)
|
||||
await asyncio.sleep(self._config.update_database_interval)
|
||||
|
||||
async def run(self, host: str, port: int) -> None:
|
||||
try:
|
||||
if self._config.log is not None:
|
||||
logger.setup_logging(self._config.log)
|
||||
if self._config.log_path is not None:
|
||||
logger.setup_logging(self._config.log_path)
|
||||
|
||||
self._logger = logger.getLogger('testdata')
|
||||
# self._logger = logger.getLogger('testdata')
|
||||
self._logger.info('Server started.')
|
||||
|
||||
coroutines = [uvicorn.Server(uvicorn.Config(self._api, host, port)).serve()]
|
||||
if self._config.database is not None:
|
||||
coroutines.append(self._update_state())
|
||||
|
||||
await asyncio.gather(*coroutines)
|
||||
await asyncio.gather(
|
||||
asyncio.create_task(uvicorn.Server(uvicorn.Config(self._api, host, port)).serve()),
|
||||
self._update_state()
|
||||
)
|
||||
except asyncio.exceptions.CancelledError:
|
||||
self._logger.info('Server stopped.')
|
||||
except Exception as err:
|
||||
|
@ -53,17 +53,6 @@ def _server(request) -> Generator[str, None, None]:
|
||||
break
|
||||
|
||||
|
||||
@pytest.mark.parametrize('_server', [({
|
||||
'keys': ['one', 'two', 'three'],
|
||||
'max-size': '100',
|
||||
'max-data': 1234,
|
||||
'buffer-size': '12MiB',
|
||||
})], indirect=['_server'])
|
||||
def test_invalid_api_key(_server):
|
||||
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=four&size=100', timeout=TIMEOUT)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.parametrize('_server', [({
|
||||
'keys': ['one', 'two', 'three'],
|
||||
'max-size': '100',
|
||||
@ -96,16 +85,13 @@ def test_request_size_upper_bound(_server):
|
||||
|
||||
@pytest.mark.parametrize('_server', [({
|
||||
'keys': ['one', 'two', 'three'],
|
||||
'max-size': '100KB',
|
||||
'max-data': '100KB',
|
||||
'max-size': '100',
|
||||
'max-data': 1234,
|
||||
'buffer-size': '12MiB',
|
||||
})], indirect=['_server'])
|
||||
def test_request_max_data_used(_server):
|
||||
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=100KB', timeout=TIMEOUT)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=1', timeout=TIMEOUT)
|
||||
assert response.status_code == 500
|
||||
def test_invalid_api_key(_server):
|
||||
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=four&size=100', timeout=TIMEOUT)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.parametrize('_server', [({
|
||||
@ -113,32 +99,18 @@ def test_request_max_data_used(_server):
|
||||
'max-size': '1KB',
|
||||
'max-data': '1KB',
|
||||
'buffer-size': '12MiB',
|
||||
'database-update-interval': 0.1
|
||||
'update-database-interval': 0.1
|
||||
})], indirect=['_server'])
|
||||
def test_check_database_update(_server):
|
||||
import importlib.metadata
|
||||
from datetime import datetime
|
||||
|
||||
database = _server
|
||||
|
||||
with open(database, 'r', encoding='utf-8') as file:
|
||||
file.seek(0)
|
||||
today = datetime.today()
|
||||
assert json.load(file) == {
|
||||
'version': importlib.metadata.version('testdata'),
|
||||
'data-used': {
|
||||
f'{today.year}-{today.month:02}': 0
|
||||
}
|
||||
}
|
||||
assert json.load(file) == {'data-used': 0}
|
||||
|
||||
response = requests.get(f'{PROTOCOL}://{HOST}:{PORT}/zeros?api_key=one&size=100', timeout=TIMEOUT)
|
||||
assert response.status_code == 200
|
||||
|
||||
time.sleep(0.1)
|
||||
file.seek(0)
|
||||
assert json.load(file) == {
|
||||
'version': importlib.metadata.version('testdata'),
|
||||
'data-used': {
|
||||
f'{today.year}-{today.month:02}': 100
|
||||
}
|
||||
}
|
||||
assert json.load(file) == {'data-used': 100}
|
||||
|
Reference in New Issue
Block a user