diff --git a/src/smolvault/config.py b/src/smolvault/config.py index ecf99c8..b0f2158 100644 --- a/src/smolvault/config.py +++ b/src/smolvault/config.py @@ -12,6 +12,7 @@ class Settings(BaseSettings): auth_secret_key: str user_whitelist: str users_limit: int + daily_upload_limit_bytes: int model_config = SettingsConfigDict(env_file=".env") diff --git a/src/smolvault/validators/operation_validator.py b/src/smolvault/validators/operation_validator.py index 69d6673..11d6317 100644 --- a/src/smolvault/validators/operation_validator.py +++ b/src/smolvault/validators/operation_validator.py @@ -4,13 +4,13 @@ from smolvault.clients.database import DatabaseClient from smolvault.config import get_settings -DAILY_UPLOAD_LIMIT_BYTES = 1_000_000_000 logger = logging.getLogger(__name__) class UploadValidator: def __init__(self) -> None: self.settings = get_settings() + self.daily_upload_limit_bytes = self.settings.daily_upload_limit_bytes self.whitelist = self.settings.user_whitelist.split(",") def upload_allowed(self, user_id: int, db_client: DatabaseClient) -> bool: @@ -24,7 +24,7 @@ def _uploads_under_limit_prev_24h(self, user_id: int, db_client: DatabaseClient) metadata = db_client.get_all_metadata(user_id, start_time=start_time) bytes_uploaded = sum([record.size for record in metadata]) logger.info("User %s has uploaded %d bytes in the last 24 hours", user_id, bytes_uploaded) - return bytes_uploaded < DAILY_UPLOAD_LIMIT_BYTES + return bytes_uploaded < self.daily_upload_limit_bytes def _user_on_whitelist(self, user_id: int) -> bool: logger.info("Checking whitelist for user %s", user_id) diff --git a/tests/testing.env b/tests/testing.env index 4020338..3cbd39a 100644 --- a/tests/testing.env +++ b/tests/testing.env @@ -4,3 +4,4 @@ SMOLVAULT_DB="test.db" SMOLVAULT_CACHE="./uploads/" USER_WHITELIST="1,2" USERS_LIMIT="3" +DAILY_UPLOAD_LIMIT_BYTES="1000000000"