From 653de5ad67eebf4a367dd5cf2c2890ba30a2e2d2 Mon Sep 17 00:00:00 2001 From: Oliver Stolpe Date: Tue, 16 Jul 2024 18:16:19 +0200 Subject: [PATCH] feat: refine quota email system (#306) (#307) --- adminsec/management/commands/import.py | 6 -- adminsec/tasks.py | 22 +++++-- adminsec/tests/test_tasks.py | 50 +++++++++++++--- config/celery.py | 8 ++- config/settings/base.py | 3 +- usersec/models.py | 83 +++++++++++++++++++------- 6 files changed, 130 insertions(+), 42 deletions(-) diff --git a/adminsec/management/commands/import.py b/adminsec/management/commands/import.py index da9cc48..6b31b2a 100644 --- a/adminsec/management/commands/import.py +++ b/adminsec/management/commands/import.py @@ -1,6 +1,5 @@ """Command for importing data from a json file.""" -import logging from collections.abc import Generator from contextlib import contextmanager @@ -37,9 +36,6 @@ def rollback(_self) -> Generator[None, None, None]: "m": "@MDC-BERLIN", } -logger = logging.getLogger(__name__) -logging.basicConfig(filename="myapp.log", level=logging.INFO) - class Command(BaseCommand): help = "Import HPC objects from a json file." @@ -55,13 +51,11 @@ def handle(self, *args, **options): try: with context, open(options["json"], "r") as jsonfile: if options["purge"]: - logger.info("Purging database of HPC objects ...") users_consented = clean_db_of_hpc_objects() if users_consented is None: self.stderr.write("Failed to clean database of HPC objects ... aborting.") return - logger.info("Importing HPC objects ...") data = HpcaccessState.model_validate_json(jsonfile.read()) for group_uuid, group_data in data.hpc_groups.items(): hpcgroup = HpcGroup( diff --git a/adminsec/tasks.py b/adminsec/tasks.py index 6dee244..f18f075 100644 --- a/adminsec/tasks.py +++ b/adminsec/tasks.py @@ -34,7 +34,7 @@ logger = logging.getLogger(__name__) -logging.basicConfig(filename="myapp.log", level=logging.INFO) +logging.basicConfig(level=logging.INFO) def _sync_ldap(write=False, verbose=False, ldapcon=None): @@ -106,19 +106,33 @@ def _generate_quota_reports(): } -@app.task(bind=True) -def send_quota_email(_self): +def _send_quota_email(status): if not settings.SEND_QUOTA_EMAILS: return reports = _generate_quota_reports() for data in reports.values(): + print(status, data) for hpc_obj, report in data.items(): - if any([not s == HpcQuotaStatus.GREEN for s in report["status"].values()]): + if any([s > status for s in report["status"].values()]): + # Skip if the object is already in a worse state + continue + + if any([s == status for s in report["status"].values()]): send_notification_storage_quota(hpc_obj, report) +@app.task(bind=True) +def send_quota_email_yellow(_self): + _send_quota_email(HpcQuotaStatus.YELLOW) + + +@app.task(bind=True) +def send_quota_email_red(_self): + _send_quota_email(HpcQuotaStatus.RED) + + @transaction.atomic @app.task(bind=True) def disable_users_without_consent(_self): diff --git a/adminsec/tests/test_tasks.py b/adminsec/tests/test_tasks.py index 3ba15b9..1edf421 100644 --- a/adminsec/tests/test_tasks.py +++ b/adminsec/tests/test_tasks.py @@ -10,10 +10,12 @@ from adminsec.ldap import LdapConnector from adminsec.tasks import ( _generate_quota_reports, + _send_quota_email, _sync_ldap, clean_db_of_hpc_objects, disable_users_without_consent, - send_quota_email, + send_quota_email_red, + send_quota_email_yellow, ) from adminsec.tests.test_ldap import ( AUTH_LDAP2_BIND_DN, @@ -43,6 +45,7 @@ HpcProjectCreateRequest, HpcProjectDeleteRequest, HpcProjectInvitation, + HpcQuotaStatus, HpcUser, HpcUserChangeRequest, HpcUserCreateRequest, @@ -244,7 +247,7 @@ def test__sync_ldap_invalid_no_user_found(self): class TestSendQuotaEmail(TestCase): - """Tests for send_quota_email.""" + """Tests for _send_quota_email.""" def setUp(self): # Superuser @@ -279,7 +282,7 @@ def setUp(self): primary_group=self.hpc_group, creator=self.user_hpcadmin, resources_requested={TIER_USER_HOME: 20}, - resources_used={TIER_USER_HOME: 10}, # 50% used + resources_used={TIER_USER_HOME: 20}, # 100% used ) self.hpc_group.owner = self.hpc_owner self.hpc_group.save() @@ -294,15 +297,15 @@ def setUp(self): primary_group=self.hpc_group, creator=self.user_hpcadmin, resources_requested={TIER_USER_HOME: 20}, - resources_used={TIER_USER_HOME: 20}, # 100% used + resources_used={TIER_USER_HOME: 18}, # 90% used ) # Create project self.hpc_project = HpcProjectFactory( group=self.hpc_group, - resources_requested={"work": 20}, - resources_used={"work": 18}, # 90% used - folders={"work": "/data/work/project"}, + resources_requested={"work": 20, "scratch": 20}, + resources_used={"work": 18, "scratch": 20}, # 90% used, 100% used + folders={"work": "/data/work/project", "scratch": "/data/scratch/project"}, ) self.hpc_project.members.add(self.hpc_owner) self.hpc_project.get_latest_version().members.add(self.hpc_owner) @@ -317,10 +320,39 @@ def test_generate_quota_reports(self): self.assertEqual(reports, expected) @override_settings(SEND_QUOTA_EMAILS=True) - def test_send_quota_email(self): - send_quota_email() + def test__send_quota_email_red(self): + _send_quota_email(HpcQuotaStatus.RED) + self.assertEqual(len(mail.outbox), 2) + + @override_settings(SEND_QUOTA_EMAILS=True) + def test__send_quota_email_red_with_delegate(self): + self.hpc_project.delegate = self.hpc_member + self.hpc_project.members.add(self.hpc_member) + self.hpc_project.get_latest_version().members.add(self.hpc_member) + self.hpc_project.save() + _send_quota_email(HpcQuotaStatus.RED) self.assertEqual(len(mail.outbox), 2) + @override_settings(SEND_QUOTA_EMAILS=True) + def test__send_quota_email_yellow(self): + _send_quota_email(HpcQuotaStatus.YELLOW) + self.assertEqual(len(mail.outbox), 1) + + @override_settings(SEND_QUOTA_EMAILS=True) + def test__send_quota_email_status_green(self): + _send_quota_email(HpcQuotaStatus.GREEN) + self.assertEqual(len(mail.outbox), 1) + + @override_settings(SEND_QUOTA_EMAILS=True) + def test_send_quota_email_red(self): + send_quota_email_red() + self.assertEqual(len(mail.outbox), 2) + + @override_settings(SEND_QUOTA_EMAILS=True) + def test_send_quota_email_yellow(self): + send_quota_email_yellow() + self.assertEqual(len(mail.outbox), 1) + class DisableUsersWithoutConsent(TestCase): """Tests for disable_users_without_consent.""" diff --git a/config/celery.py b/config/celery.py index 3a8be87..ea2bb21 100644 --- a/config/celery.py +++ b/config/celery.py @@ -20,8 +20,12 @@ "schedule": crontab(hour=0, minute=5), # "args": (True, False), }, - "send_quota_email": { - "task": "adminsec.tasks.send_quota_email", + "send_quota_email_yellow": { + "task": "adminsec.tasks.send_quota_email_yellow", + "schedule": crontab(day_of_week=1, hour=0, minute=20), + }, + "send_quota_email_red": { + "task": "adminsec.tasks.send_quota_email_red", "schedule": crontab(hour=0, minute=10), }, "disable_users_without_consent": { diff --git a/config/settings/base.py b/config/settings/base.py index 7b11c59..b570fce 100644 --- a/config/settings/base.py +++ b/config/settings/base.py @@ -412,7 +412,8 @@ EMAIL_PORT = env.int("EMAIL_PORT", 25) EMAIL_SENDER = env.str("EMAIL_SENDER", "root@admin") SEND_QUOTA_EMAILS = env.bool("SEND_QUOTA_EMAILS", False) -QUOTA_WARNING_THRESHOLD = env.int("QUOTA_WARNING_THRESHOLD", 80) +QUOTA_WARNING_THRESHOLD = env.int("QUOTA_WARNING_THRESHOLD", 90) +QUOTA_WARNING_ABSOLUTE = env.int("QUOTA_WARNING_ABSOLUTE", 5) CONSENT_GRACE_PERIOD = env.int("CONSENT_GRACE_PERIOD", 30) VIEW_MODE = env.bool("VIEW_MODE", False) diff --git a/usersec/models.py b/usersec/models.py index 3a58a38..63f9f7b 100644 --- a/usersec/models.py +++ b/usersec/models.py @@ -273,9 +273,21 @@ def get_retract_url(self): @unique class HpcQuotaStatus(Enum): - RED = "red" - YELLOW = "yellow" - GREEN = "green" + GREEN = 1 + YELLOW = 2 + RED = 3 + + def __ge__(self, other): + return self.value >= other.value + + def __gt__(self, other): + return self.value > other.value + + def __le__(self, other): + return self.value <= other.value + + def __lt__(self, other): + return self.value < other.value class CheckQuotaMixin: @@ -322,7 +334,10 @@ def generate_quota_report(self): if used_val >= requested_val: result["status"][key] = HpcQuotaStatus.RED - elif used_val >= requested_val * (settings.QUOTA_WARNING_THRESHOLD / 100): + elif used_val >= max( + requested_val * settings.QUOTA_WARNING_THRESHOLD / 100, + requested_val - settings.QUOTA_WARNING_ABSOLUTE, + ): result["status"][key] = HpcQuotaStatus.YELLOW else: @@ -616,19 +631,33 @@ def __str__(self): self_delegate_username, ) - def get_manager_emails(self): - emails = [self.owner.user.email] + def get_manager_emails(self, slim=True): + owner_email = self.owner.user.email + delegate_email = self.delegate.user.email if self.delegate else None - if self.delegate: - emails.append(self.delegate.user.email) + if slim: + emails = [delegate_email if self.delegate else owner_email] + + else: + emails = [owner_email] + + if self.delegate: + emails.append(delegate_email) return emails - def get_manager_names(self): - names = [self.owner.user.get_full_name()] + def get_manager_names(self, slim=True): + owner_name = self.owner.user.get_full_name() + delegate_name = self.delegate.user.get_full_name() if self.delegate else None - if self.delegate: - names.append(self.delegate.user.get_full_name()) + if slim: + names = [delegate_name if self.delegate else owner_name] + + else: + names = [owner_name] + + if self.delegate: + names.append(delegate_name) return names @@ -793,19 +822,33 @@ def __str__(self): self_delegate_username, ) - def get_manager_emails(self): - emails = [self.group.owner.user.email] + def get_manager_emails(self, slim=True): + owner_email = self.group.owner.user.email + delegate_email = self.delegate.user.email if self.delegate else None - if self.delegate: - emails += self.delegate.user.email + if slim: + emails = [delegate_email if self.delegate else owner_email] + + else: + emails = [owner_email] + + if self.delegate: + emails.append(delegate_email) return emails - def get_manager_names(self): - names = [self.group.owner.user.get_full_name()] + def get_manager_names(self, slim=True): + owner_name = self.group.owner.user.get_full_name() + delegate_name = self.delegate.user.get_full_name() if self.delegate else None - if self.delegate: - names += self.delegate.user.get_full_name() + if slim: + names = [delegate_name if self.delegate else owner_name] + + else: + names = [owner_name] + + if self.delegate: + names.append(delegate_name) return names