attempted changes to create test users and applications

Author: clewellyn-nava

apps/testclient/management/commands/create_test_users_and_applications_batch.py

@@ -5,7 +5,7 @@
 
 import apps.logging.request_logger as logging
 
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from os import listdir
 from random import randint, randrange, choice, sample
 from django.contrib.auth.models import Group
@@ -15,6 +15,7 @@
 from django.conf import settings
 from django.utils import timezone
 from oauth2_provider.models import AccessToken, RefreshToken
+from tqdm import tqdm
 
 from apps.accounts.models import UserProfile
 from apps.fhir.bluebutton.models import Crosswalk
@@ -34,7 +35,9 @@
                            'capability-a', 'capability-b']
 APPLICATION_SCOPES_NON_DEMOGRAPHIC = ['patient/ExplanationOfBenefit.read',
                                       'patient/Coverage.read', 'capability-a', 'capability-b']
-DEFAULT_BENE_COUNT = 100
+
+# Keep up with prod count
+DEFAULT_BENE_COUNT = 675000
 DEFAULT_DEV_COUNT = 150
 DEFAULT_MAX_APPS_PER_DEV = 5
 
@@ -60,92 +63,114 @@ def get_first_available_number(firstname):
     return int(begin) + 1
 
 def create_dev_users_apps_and_bene_crosswalks(
-        group,
         bene_count,
         dev_count,
         app_max,
         refresh_count,
         archived_tokens,
-        archived_grants
+        archived_grants,
+        auto_generate
 ):
-    #
-    # generate dev users dev0001, dev0002, dev0003 etc, with password, email, security questions, etc.
-    # Counts are based on the inputs, but if undefined,will default to the defaults defined above
-    # each dev user can have 1 to many applications: dev0001-app01, dev0001-app02
-    # generate crosswalk bene users with FHIR-ID. HICN-HASH, MBI-HASH etc
-    #
-
-    # effectively turn off logging
+    """Creates bene and dev users, associated apps, and crosswalks for testing
+
+    Args:
+        bene_count (_type_): _description_
+        dev_count (_type_): _description_
+        app_max (_type_): _description_
+        refresh_count (_type_): _description_
+        archived_tokens (_type_): _description_
+        archived_grants (_type_): _description_
+        auto_generate (_type_): _description_
+    """
     mymedicare_cb_logger.setLevel(logging.CRITICAL)
     outreach_logger.setLevel(logging.CRITICAL)
 
+    group = create_group()
     bene_pk_list = []
-    files = [f for f in listdir('./synthetic-data') if re.match(r'synthetic-beneficiary-.*\.rif', f)]
+    
     count = 0
     file_cnt = 0
     synthetic_bene_cnt = 0
     begin_number = get_first_available_number('bene')
-
-    for f in files:
-        print('file={}'.format(f))
-        bene_rif = open('./synthetic-data/{}'.format(f), 'r')
-        while True:
-            if count == bene_count:
-                break
-            line = bene_rif.readline()
-            if not line:
+    if auto_generate:
+        for i in tqdm(range(bene_count)):
+            id = str(i).zfill(10)
+            fn = f'FN{id}'
+            ln = f'LN{id}'
+            fhir_id = f'-{id}'
+            args = {
+                'username': str(uuid.uuid4()),
+                'user_hicn_hash': hash_hicn(id),
+                'user_mbi': f'MBI{id}',
+                'user_id_type': 'H',
+                'fhir_id_v2': fhir_id,
+                'first_name': fn,
+                'last_name': ln,
+                'email': f'{fn}.{ln}@xyz.com'
+            }
+            slsx_client = OAuth2ConfigSLSx(args)
+            user = create_beneficiary_record(slsx_client, fhir_id)
+            bene_pk_list.append(user.pk)
+    else:
+        files = [f for f in listdir('./synthetic-data') if re.match(r'synthetic-beneficiary-.*\.rif', f)]
+        for f in files:
+            print('file={}'.format(f))
+            bene_rif = open('./synthetic-data/{}'.format(f), 'r')
+            while True:
+                if count == bene_count:
+                    break
+                line = bene_rif.readline()
+                if not line:
+                    break
+                if line.startswith('DML_IND|BENE_ID'):
+                    continue
+                else:
+                    flds = line.split('|')
+                    fhir_id = flds[1]
+                    mbi = flds[14]
+                    hicn = flds[18]
+                    fn = f'bene{count + begin_number}'
+                    ln = f'user{count + begin_number}'
+
+                    # skip fred
+                    if fhir_id != '-20140000008325':
+                        args = {
+                            'username': str(uuid.uuid1()),
+                            'user_hicn_hash': hash_hicn(hicn),
+                            'user_mbi': mbi,
+                            'user_id_type': 'H',
+                            'fhir_id_v2': fhir_id,
+                            'first_name': fn,
+                            'last_name': ln,
+                            'email': fn + '.' + ln + '@xyz.net',
+                        }
+                        slsx_client = OAuth2ConfigSLSx(args)
+                        try:
+                            u = create_beneficiary_record(slsx_client, fhir_id)
+                            date_picked = datetime.now(timezone.utc) - timedelta(days=randrange(700))
+                            u.date_joined = date_picked.replace(tzinfo=pytz.utc)
+                            u.save()
+                            c = Crosswalk.objects.get(user=u)
+                            c.created = u.date_joined
+                            c.save()
+                            bene_pk_list.append(u.pk)
+                            synthetic_bene_cnt += 1
+                            count += 1
+                        except ValidationError:
+                            # If there is something wrong during 'create_beneficiary_record'
+                            # i.e. 'user already exists', just try the next .rif record
+                            continue
+            bene_rif.close()
+            file_cnt += 1
+            print('RIF file processed = {}, synthetic bene generated = {}'.format(file_cnt, synthetic_bene_cnt))
+            if file_cnt >= 1:
                 break
-            if line.startswith('DML_IND|BENE_ID'):
-                continue
-            else:
-                flds = line.split('|')
-                fhir_id = flds[1]
-                mbi = flds[14]
-                hicn = flds[18]
-                fn = 'bene{}'.format(count + begin_number)
-                ln = 'user{}'.format(count + begin_number)
-
-                # skip fred
-                if fhir_id != '-20140000008325':
-                    args = {
-                        'username': str(uuid.uuid1()),
-                        'user_hicn_hash': hash_hicn(hicn),
-                        'user_mbi': mbi,
-                        'user_id_type': 'H',
-                        'fhir_id_v2': fhir_id,
-                        'first_name': fn,
-                        'last_name': ln,
-                        'email': fn + '.' + ln + '@xyz.net',
-                    }
-                    slsx_client = OAuth2ConfigSLSx(args)
-                    try:
-                        u = create_beneficiary_record(slsx_client, fhir_id)
-                        date_picked = datetime.utcnow() - timedelta(days=randrange(700))
-                        u.date_joined = date_picked.replace(tzinfo=pytz.utc)
-                        u.save()
-                        c = Crosswalk.objects.get(user=u)
-                        c.created = u.date_joined
-                        c.save()
-                        bene_pk_list.append(u.pk)
-                        synthetic_bene_cnt += 1
-                        count += 1
-                        print('.', end='', flush=True)
-                        time.sleep(.05)
-                    except ValidationError:
-                        # If there is something wrong during 'create_beneficiary_record'
-                        # i.e. 'user already exists', just try the next .rif record
-                        continue
-        bene_rif.close()
-        file_cnt += 1
-        print('RIF file processed = {}, synthetic bene generated = {}'.format(file_cnt, synthetic_bene_cnt))
-        if file_cnt >= 1:
-            break
 
     # create dev users according dev-count parameter, default to 100
     # generate access tokens + refresh tokens + archived tokens for random picked benes for each app
     app_index = 0
     begin_number = get_first_available_number('DevUserFN')
-    for i in range(dev_count):
+    for i in tqdm(range(dev_count)):
         dev_u_fn = 'DevUserFN{}'.format(i + begin_number)
         dev_u_ln = 'DevUserLN{}'.format(i + begin_number)
         u = User.objects.create_user(username='{}.{}'.format(dev_u_fn, dev_u_ln),
@@ -166,7 +191,6 @@ def create_dev_users_apps_and_bene_crosswalks(
         u.groups.add(group)
         # apps per DEV user
         app_cnt = randint(1, app_max) if app_max > 0 else 0
-        print('>>>>generating apps for user={}'.format(u.username))
         for i in range(app_cnt):
             app_index += 1
             app_name = 'app{}_{}'.format(i, u)
@@ -205,7 +229,6 @@ def create_dev_users_apps_and_bene_crosswalks(
             for t in titles:
                 c = ProtectedCapability.objects.get(title=t)
                 a.scope.add(c)
-        print('<<<<<generated apps for user={}'.format(u.username))
 
     # go through benes: each bene sign up to 1, 2, 3 apps
     # most 70% 1 app, 25% 2 apps, 5% 3 apps
@@ -274,14 +297,12 @@ def create_test_access_refresh_archived_objects(
                                     scope=scope)
     at.created = date_created.replace(tzinfo=pytz.utc)
     at.save()
-    print('<<< access token created for ' + user.username + '/' + application.name)
 
     for i in range(refresh_count):
         rt = RefreshToken.objects.create(user=user, application=application,
                                      token=uuid.uuid4().hex)
         rt.created = at.created
         rt.save()
-        print('<<< ' + user.username + ' refresh token ' + str(i) + ' generated')
 
     # archived token: created, updated, archived_at datetime fields
     for i in range(archived_token_count):
@@ -297,7 +318,6 @@ def create_test_access_refresh_archived_objects(
         date_archived = ot.created + timedelta(days=10)
         ot.archived_at = date_archived.replace(tzinfo=pytz.utc)
         ot.save()
-        print('<<< ' + user.username + ' archived token ' + str(i) + ' generated')
 
     past_date = timezone.now() - timedelta(days=2)
     for i in range(archived_grant_count):
@@ -308,8 +328,11 @@ def create_test_access_refresh_archived_objects(
                                       archived_at=past_date)
         past_date = past_date - timedelta(days=2)
         adag.save()
-        print('<<< ' + user.username + 'archived grant ' + str(i) + ' generated')
 
+def delete_users_apps_crosswalks():
+    User.objects.all().delete()
+    Application.objects.all().delete()
+    Crosswalk.objects.all().delete()
 
 class Command(BaseCommand):
     help = ('Create dev users and create'
@@ -333,6 +356,8 @@ def add_arguments(self, parser):
         parser.add_argument('-g', '--archived-access-grants', default=0,
                             help='Archived access grants per user/app combination. '
                             'If none, defaults to 0.')
+        parser.add_argument('--auto-generate', default=False, help='Do not prompt for user input.')
+        parser.add_argument('--delete', default=False, help='Delete all test users and applications.')
 
     def handle(self, *args, **options):
         bene_count = int(options['bene_count'])
@@ -341,14 +366,18 @@ def handle(self, *args, **options):
         refresh_tokens = int(options['refresh_tokens'])
         archived_tokens = int(options['archived_tokens'])
         archived_grants = int(options['archived_access_grants'])
-        g = create_group()
+        auto_generate = bool(options['auto_generate'])
+        delete = bool(options['delete'])
+        
+        if delete:
+            delete_users_apps_crosswalks()
         create_dev_users_apps_and_bene_crosswalks(
-            g,
             bene_count,
             dev_count,
             app_max,
             refresh_tokens,
             archived_tokens,
-            archived_grants)
+            archived_grants,
+            auto_generate)
         # update grants
         update_grants()

apps/testclient/management/commands/load_test_data.py

@@ -0,0 +1,39 @@
+from django.core.management.base import BaseCommand
+from django.contrib.auth.models import User
+from oauth2_provider.models import AccessToken, RefreshToken
+from apps.fhir.bluebutton.models import Crosswalk
+from apps.dot_ext.models import ArchivedToken
+from apps.capabilities.models import ProtectedCapability
+from apps.authorization.models import update_grants, ArchivedDataAccessGrant
+from oauth2_provider.models import get_application_model
+
+
+def flush_database():
+    """Flush all data from the database."""
+    print('Flushing database...')
+    User.objects.all()
+    get_application_model().objects.all()
+    Crosswalk.objects.all()
+    ProtectedCapability.objects.all()
+    AccessToken.objects.all()
+    RefreshToken.objects.all()
+    ArchivedToken.objects.all()
+    ArchivedDataAccessGrant.objects.all()
+    print('Database flushed.')
+
+class Command(BaseCommand):
+    help = ('Create dev users and create'
+            ' apps for each of them, create bene users from '
+            'synthetic data and crosswalk for each bene.')
+
+    def add_arguments(self, parser):
+        parser.add_argument('--flush', default=False, help='Flush the database')
+
+    def handle(self, *args, **options):
+        delete = bool(options['delete'])
+        
+        if delete:
+            flush_database()
+        
+        # update grants
+        update_grants()

requirements/requirements.dev.in

@@ -11,5 +11,7 @@ selenium
 pytest 
 Flask 
 Flask-WTF
+typing_extensions>=4.6.0
 setuptools_scm<7,>=3.2
 setuptools>=40.8.0
+tqdm

requirements/requirements.dev.txt

@@ -866,6 +866,10 @@ tomli==2.2.1 \
     --hash=sha256:ece47d672db52ac607a3d9599a9d48dcb2f2f735c6c2d1f34130085bb12b112a \
     --hash=sha256:f4039b9cbc3048b2416cc57ab3bda989a6fcf9b36cf8937f01a6e731b64f80d7
     # via setuptools-scm
+tqdm==4.67.1 \
+    --hash=sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2 \
+    --hash=sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2
+    # via -r requirements/requirements.dev.in
 trio==0.24.0 \
     --hash=sha256:c3bd3a4e3e3025cd9a2241eae75637c43fe0b9e88b4c97b9161a55b9e54cd72c \
     --hash=sha256:ffa09a74a6bf81b84f8613909fb0beaee84757450183a7a2e0b47b455c0cac5d
@@ -876,10 +880,11 @@ trio-websocket==0.11.1 \
     --hash=sha256:18c11793647703c158b1f6e62de638acada927344d534e3c7628eedcb746839f \
     --hash=sha256:520d046b0d030cf970b8b2b2e00c4c2245b3807853ecd44214acd33d74581638
     # via selenium
-typing-extensions==4.5.0 \
-    --hash=sha256:5cb5f4a79139d699607b3ef622a1dedafa84e115ab0024e0d9c044a9479ca7cb \
-    --hash=sha256:fb33085c39dd998ac16d1431ebc293a8b3eedd00fd4a32de0ff79002c19511b4
+typing-extensions==4.15.0 \
+    --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \
+    --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548
     # via
+    #   -r requirements/requirements.dev.in
     #   -r requirements/requirements.in
     #   dj-database-url
     #   jwcrypto

requirements/requirements.in

@@ -37,7 +37,7 @@ pyyaml
 py
 attrs
 pyrsistent
-typing_extensions
+typing_extensions>=4.6.0
 importlib-metadata
 certifi
 

requirements/requirements.txt

@@ -638,9 +638,9 @@ sqlparse==0.5.0 \
     # via
     #   -r requirements/requirements.in
     #   django
-typing-extensions==4.5.0 \
-    --hash=sha256:5cb5f4a79139d699607b3ef622a1dedafa84e115ab0024e0d9c044a9479ca7cb \
-    --hash=sha256:fb33085c39dd998ac16d1431ebc293a8b3eedd00fd4a32de0ff79002c19511b4
+typing-extensions==4.15.0 \
+    --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \
+    --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548
     # via
     #   -r requirements/requirements.in
     #   dj-database-url

synthetic-data/synthetic-beneficiary.csv

@@ -0,0 +1 @@
+BENE_HICN, BENE_MBI, BENE_FHIR, BENE_FN, BENE_LN