Added an option to insert attachemnts onto the loads with the fake data generator

2 files changed, 40 insertions, 8 deletions

diff --git a/app/chicken.pdf b/app/chicken.pdf
new file mode 100644
index 0000000..09ac530
--- /dev/null
+++ b/app/chicken.pdf
diff --git a/app/dispatch/management/commands/insert_fake_data.py b/app/dispatch/management/commands/insert_fake_data.py
index d7169fe..3800372 100644
--- a/app/dispatch/management/commands/insert_fake_data.py
+++ b/app/dispatch/management/commands/insert_fake_data.py
@@ -1,6 +1,7 @@
 from django.core.management.base import BaseCommand, CommandError
+from django.core.files import File
 from django.contrib.auth import get_user_model
-from dispatch.models import Customer, Load, Identity
+from dispatch.models import Customer, Load, Identity, Paperwork
 from faker import Faker
 import random
 
@@ -15,7 +16,7 @@ class Command(BaseCommand):
     help = """Generate fake data.\n
     Example usage: \n
    ./manage.py insert_fake_data --companies 4 --users 3 --loads 800
-   --start-date='-16w' --end-date '+16w' """
+   --start-date='-16w' --end-date '+16w' --attachments=True"""
 
     fake = Faker()
     end_date = ""
@@ -24,6 +25,7 @@ class Command(BaseCommand):
     def add_arguments(self, parser):
         # parser.add_argument('--file', type=str, dest='filename')
         parser.add_argument('--companies', type=int, dest='companies', default='0')
+        parser.add_argument('--attachments', type=bool, dest='attachments', default=False)
         parser.add_argument('--users', type=int, dest='users', default='0')
         parser.add_argument('--loads', type=int, dest='loads', default='0')
         parser.add_argument('--start-date', type=str, dest='startdate', default='-2w')
@@ -35,20 +37,27 @@ class Command(BaseCommand):
         self.end_date   = options['enddate']
         self.start_date = options['startdate']
 
-        customer_ids = []
-        user_ids = []
+        # customer_ids = []
+        # user_ids = []
 
         for i in range(0, options['companies']):
             co = self.fake_customer()
 
             # We're going to hold the model IDs in memory to access randomly
             # later on
-            customer_ids.append(co.pk)
+            # customer_ids.append(co.pk)
 
         for i in range(0, options['users']):
             u = self.fake_user()
 
-            user_ids.append(u.pk)
+            # user_ids.append(u.pk)
+
+
+        # We're going to go about getting the list of user and customer IDs
+        # in a slightly different manner now so that we get _all_ customers
+        # and _all_ users.
+        customer_ids = self.get_customer_ids()
+        user_ids = self.get_user_ids()
 
         for i in range(0, options['loads']):
             co_id = customer_ids[random.randint(0, len(customer_ids)-1)]
@@ -57,7 +66,7 @@ class Command(BaseCommand):
             u_id = user_ids[random.randint(0, len(user_ids)-1)]
             u = User.objects.get(pk=u_id)
 
-            l = self.fake_load(u, co)
+            l = self.fake_load(u, co, options['attachments'])
 
     def fake_customer(self):
         new_customer = Customer(
@@ -70,7 +79,7 @@ class Command(BaseCommand):
         new_customer.save()
         return new_customer
 
-    def fake_load(self, usr, co):
+    def fake_load(self, usr, co, attachments=False):
         fake_description = "{}-{} {} {}K".format(
             self.fake.month(), self.fake.day_of_month(),
             self.fake.license_plate(), random.randint(5, 50))
@@ -87,8 +96,18 @@ class Command(BaseCommand):
         # Set the date to something random between last week and next week
         new_load.date=self.fake.date_between(start_date=self.start_date, end_date=self.end_date)
         new_load.save()
+
+        if attachments:
+            print("yeppers")
+            self.add_attachment(new_load)
+
         return new_load
 
+    def add_attachment(self, load):
+        p = Paperwork(load=load)
+        p.document = File(open('chicken.pdf', 'rb'))
+        p.save()
+
 
     def fake_user(self):
         ffname = self.fake.first_name()
@@ -117,3 +136,16 @@ class Command(BaseCommand):
         except IntegrityError:
             # Around and around we go until we get something new
             return self.fake_user()
+
+
+    def get_customer_ids(self):
+        id_array = []
+        for cus in Customer.objects.all():
+            id_array.append(cus.pk)
+        return id_array
+
+    def get_user_ids(self):
+        id_array = []
+        for u in User.objects.all():
+            id_array.append(u.pk)
+        return id_array