From bc1086e37b4f1d724ac44c2df4d3c3a9e88f3cd3 Mon Sep 17 00:00:00 2001 From: Steve Astels Date: Thu, 1 Aug 2024 15:47:15 -0400 Subject: [PATCH] Add command to generate test data (#2235) --- app/commands/test_data.py | 172 +++++++++++++++++++++++++++++++++++--- 1 file changed, 160 insertions(+), 12 deletions(-) diff --git a/app/commands/test_data.py b/app/commands/test_data.py index 929c3a006a..89cc04e47e 100644 --- a/app/commands/test_data.py +++ b/app/commands/test_data.py @@ -1,15 +1,26 @@ import functools +import math +import random import uuid +from datetime import datetime import click +from dateutil import rrule from flask import cli as flask_cli +from app import db from app.dao.services_dao import ( + dao_create_service, dao_fetch_all_services_by_user, delete_service_and_all_associated_db_objects, ) -from app.dao.users_dao import delete_model_user, delete_user_verify_codes -from app.models import User +from app.dao.templates_dao import dao_create_template +from app.dao.users_dao import ( + delete_model_user, + delete_user_verify_codes, + save_model_user, +) +from app.models import NotificationHistory, Organisation, Service, Template, User @click.group(name="test-data", help="Generate and destroy test data") @@ -42,32 +53,169 @@ def setup_test_data_commands(application): @test_data_command() @click.option( - "-u", - "--user_email_prefix", - required=True, + "-p", + "--prefix", + default="notify-test-data", + show_default=True, help=""" - Functional test user email prefix. eg "notify-test-preview" + test data prefix" """, ) # noqa -def purge_functional_test_data(user_email_prefix): +@click.option("-s", "--num_services", default=1, show_default=True, help="Number of services to create") +@click.option("-n", "--num_notifications", default=1, show_default=True, help="Number of notifications to create") +@click.option("-b", "--batch_size", default=10000, show_default=True, help="Number of notifications to create in each batch") +def generate(prefix, num_services, num_notifications, batch_size): + """ + Generate test data """ - Remove non-seeded functional test data + print("Building org...") + org = Organisation( + name=f"{prefix} {uuid.uuid4()}", + organisation_type="central", + ) + db.session.add(org) + db.session.flush() + print(" -> Done.") + + print(f"Building {num_services} services...") + services = [] + templates = [] + + for batch in range(num_services): + data_prefix = f"{prefix}+{uuid.uuid4()}" + user_email = f"{data_prefix}@cds-snc.ca" + data = { + "id": uuid.uuid4(), + "name": user_email, + "email_address": user_email, + "password": f"{uuid.uuid4()}", + "mobile_number": "16135550123", + "state": "active", + "blocked": False, + } + user = User(**data) + save_model_user(user) + + service = Service( + organisation_id=org.id, + name=f"{data_prefix} service {batch}", + created_by_id=user.id, + active=True, + restricted=False, + organisation_type="central", + message_limit=250_000, + sms_daily_limit=10_000, + email_from=f"{data_prefix}_{batch}@notify.works", + ) + services.append(service) + dao_create_service(service, user) + + service_templates = { + "email": Template( + name=f"{data_prefix}: email", + service_id=service.id, + template_type="email", + subject="email", + content="email body", + created_by_id=user.id, + ), + "sms": Template( + name=f"{data_prefix}: sms", + service_id=service.id, + template_type="sms", + subject="sms", + content="sms body", + created_by_id=user.id, + ), + } + dao_create_template(service_templates["email"]) + dao_create_template(service_templates["sms"]) + templates.append(service_templates) + db.session.flush() + print(" -> Done.") + + num_batches = math.ceil(num_notifications / batch_size) + print(f"Building {num_notifications} notifications in batches of {batch_size}...") + last_new_year = datetime(datetime.today().year - 1, 1, 1, 12, 0, 0) + daily_dates_since_last_new_year = list(rrule.rrule(freq=rrule.DAILY, dtstart=last_new_year, until=datetime.today())) + for batch in range(num_batches): + print(f" -> Building batch #{batch + 1} of {num_batches}...") + notifications_batch = [] + for _ in range(min(batch_size, num_notifications - (batch * batch_size))): + notification_type = random.choice(["sms", "email"]) + service_index = random.choice(range(len(services))) + service = services[service_index] + template = templates[service_index][notification_type] + notifications_batch.append( + NotificationHistory( + id=uuid.uuid4(), + job_id=None, + job_row_number=None, + service_id=service.id, + template_id=template.id, + template_version=1, + api_key_id=None, + key_type="normal", + billable_units=1, + rate_multiplier=1, + notification_type=notification_type, + created_at=random.choice(daily_dates_since_last_new_year), + status="delivered", + client_reference=data_prefix, + ) + ) + print(f" -> Adding {len(notifications_batch)} notifications...") + db.session.bulk_save_objects(notifications_batch) + db.session.flush() + print(" -> Done.") + + print("Committing...") + db.session.commit() + print("Finished.") + - users, services, etc. Give an email prefix. Probably "notify-test-preview". +@test_data_command() +@click.option( + "-p", + "--prefix", + default="notify-test-data", + show_default=True, + help=""" + test data prefix" +""", +) # noqa +def delete(prefix): """ - users = User.query.filter(User.email_address.like("{}%".format(user_email_prefix))).all() + Delete test data from the database + """ + print(f"\n\nDeleting test data for prefix {prefix}...\n") + users = User.query.filter(User.email_address.like("{}%".format(prefix))).all() + org_ids = [] + if len(users) == 0: + print("No users found with email prefix {}".format(prefix)) + return for usr in users: - # Make sure the full email includes a uuid in it - # Just in case someone decides to use a similar email address. try: uuid.UUID(usr.email_address.split("@")[0].split("+")[1]) except ValueError: print("Skipping {} as the user email doesn't contain a UUID.".format(usr.email_address)) else: + print(f"Deleting user {usr.email_address} and related data...") services = dao_fetch_all_services_by_user(usr.id) if services: for service in services: + org_ids.append(service.organisation_id) delete_service_and_all_associated_db_objects(service) else: delete_user_verify_codes(usr) delete_model_user(usr) + print(" -> Done.") + + db.session.commit() + + for org_id in set(org_ids): + print(f"Deleting organisation {org_id}...") + Organisation.query.filter_by(id=org_id).delete(synchronize_session=False) + db.session.commit() + + print("Finished.")