From 06eb1b66831a3e27068dafbd596dda45bdbece53 Mon Sep 17 00:00:00 2001
From: AlanCoding <arominge@redhat.com>
Date: Thu, 20 Feb 2020 20:53:11 -0500
Subject: [PATCH] Randomize JT for each job batch

Populate some more fields just to have them populated

Include some missing ForeignKey links for data integrity

Add some more to fields from JT set
---
 tools/scripts/firehose.py | 58 ++++++++++++++++++++++++++++-----------
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/tools/scripts/firehose.py b/tools/scripts/firehose.py
index 58f3d68b43..86e1d00a8c 100755
--- a/tools/scripts/firehose.py
+++ b/tools/scripts/firehose.py
@@ -57,6 +57,8 @@ dsn = f'dbname={name} user={user} password={pw} host={host}'
 
 u = str(uuid4())
 
+STATUS_OPTIONS = ('successful', 'failed', 'error', 'canceled')
+
 
 class YieldedRows(StringIO):
 
@@ -129,17 +131,35 @@ def generate_jobs(jobs):
     from awx.main.models import UnifiedJob, Job, JobTemplate
     fields = list(set(Job._meta.fields) - set(UnifiedJob._meta.fields))
     job_field_names = set([f.attname for f in fields])
-    jt = JobTemplate.objects.first()
-    jt_defaults = dict(
-        (f.attname, getattr(jt, f.attname))
-        for f in JobTemplate._meta.get_fields()
-        if f.editable and f.attname in job_field_names and getattr(jt, f.attname)
-    )
-    jt_defaults['job_template_id'] = jt.pk
+    # extra unified job field names from base class
+    for field_name in ('name', 'created_by_id', 'modified_by_id'):
+        job_field_names.add(field_name)
+    jt_count = JobTemplate.objects.count()
+
+    def make_batch(N, jt_pos=0):
+        jt = None
+        while not jt:
+            try:
+                jt = JobTemplate.objects.all()[jt_pos % jt_count]
+            except IndexError as e:
+                # seems to happen every now and then due to some race condition
+                print('Warning: IndexError on {} JT, error: {}'.format(
+                    jt_pos % jt_count, e
+                ))
+            jt_pos += 1
+        jt_defaults = dict(
+            (f.attname, getattr(jt, f.attname))
+            for f in JobTemplate._meta.get_fields()
+            if f.editable and f.attname in job_field_names and getattr(jt, f.attname)
+        )
+        jt_defaults['job_template_id'] = jt.pk
+        jt_defaults['unified_job_template_id'] = jt.pk  # populated by save method
 
-    def make_batch(N, **extra):
         jobs = [
-            Job(status='canceled', created=now(), modified=now(), elapsed=0., **extra)
+            Job(
+                status=STATUS_OPTIONS[i % len(STATUS_OPTIONS)],
+                started=now(), created=now(), modified=now(), finished=now(),
+                elapsed=0., **jt_defaults)
             for i in range(N)
         ]
         ujs = UnifiedJob.objects.bulk_create(jobs)
@@ -148,14 +168,16 @@ def generate_jobs(jobs):
         with connection.cursor() as cursor:
             query, params = query.sql_with_params()[0]
             cursor.execute(query, params)
-        return ujs[-1]
+        return ujs[-1], jt_pos
 
     i = 1
+    jt_pos = 0
+    s = time()
     while jobs > 0:
-        s = time()
+        s_loop = time()
         print('running batch {}, runtime {}'.format(i, time() - s))
-        created = make_batch(min(jobs, 1000), **jt_defaults)
-        print('took {}'.format(time() - s))
+        created, jt_pos = make_batch(min(jobs, 1000), jt_pos)
+        print('took {}'.format(time() - s_loop))
         i += 1
         jobs -= 1000
     return created
@@ -254,9 +276,13 @@ def generate_events(events, job):
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--jobs', type=int, default=1000000) # 1M by default
-    parser.add_argument('--events', type=int, default=1000000000) # 1B by default
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        '--jobs', type=int, help='Number of jobs to create.',
+        default=1000000) # 1M by default
+    parser.add_argument(
+        '--events', type=int, help='Number of events to create.',
+        default=1000000000) # 1B by default
     params = parser.parse_args()
     jobs = params.jobs
     events = params.events