From 3063bc0167f198fa4c4dd04711179a0bf0fe7b81 Mon Sep 17 00:00:00 2001
From: Kristian Krsnik <git@krsnik.at>
Date: Sat, 3 Aug 2024 01:05:22 +0200
Subject: [PATCH] created a type model for the config file

---
 README.md                |  2 +-
 testdata/custom_types.py | 33 +++++++++++++++------------------
 testdata/main.py         | 14 +++-----------
 3 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index e9b9240..dc05fc8 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@
     "log": "-",
     "buffer-size": "4KiB",
     "max-size": "2GB",
-    "keys": [
+    "api-keys": [
         "TESTKEY"
     ],
     "max-data": "10GB",
diff --git a/testdata/custom_types.py b/testdata/custom_types.py
index 142624c..83588bf 100644
--- a/testdata/custom_types.py
+++ b/testdata/custom_types.py
@@ -1,27 +1,24 @@
-from pydantic import BaseModel, ConfigDict, field_validator
+from pydantic import BaseModel, Field, 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')
+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 = '-'
 
-#     @field_validator('buffer_size', 'max_size', 'max_data')
-#     @classmethod
-#     def convert_size(cls, value: int | str) -> int:
-#         return convert_to_bytes(value)
+    model_config = ConfigDict(extra='forbid')
 
-
-# TODO
-# class DataBase(BaseModel):
-#     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)
 
 
 class TestDataBody(BaseModel):
diff --git a/testdata/main.py b/testdata/main.py
index 62b57f8..f435580 100644
--- a/testdata/main.py
+++ b/testdata/main.py
@@ -8,6 +8,7 @@ import uvicorn
 
 from .utils import convert_to_bytes, save_database
 from .api import create_api
+from .custom_types import Config
 
 
 # Setup Parser
@@ -35,14 +36,5 @@ def run(host: str, port: int, api_keys: set[str], max_size: int, max_data: int,
 
 def main():
     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'])
-    max_size = convert_to_bytes(config['max-size'])
-    max_data = convert_to_bytes(config['max-data'])
-    api_keys = config['keys']
-    database = config['database']
-
-    run(host, port, api_keys, max_size, max_data, database, buffer_size)
+    config = Config.model_validate_json(args.config.read())
+    run(**config.model_dump(exclude={'log'}))