1 files changed, 26 insertions, 5 deletions
diff --git a/app/dispatch/views.py b/app/dispatch/views.py
index ac399f6..bf4c3fb 100644
--- a/app/dispatch/views.py
+++ b/app/dispatch/views.py
@@ -10,7 +10,7 @@ from django.views.generic import TemplateView,ListView
 from django.views.generic.detail import DetailView
 from django.views.generic.edit import CreateView, UpdateView, DeleteView
 from django.core.urlresolvers import reverse_lazy
-from dispatch.models import Customer, Load, Paperwork, Invoice, Settings, Identity
+from dispatch.models import Customer, Load, Paperwork, Invoice, UserInvoiceNumber, Settings, Identity
 from dispatch.forms import AddPaperworkForm
 from django.contrib.auth.models import User
 from django.contrib.auth.mixins import UserPassesTestMixin
@@ -307,15 +307,36 @@ def PaperworkDownload(request, load_id, pk):
     except Exception as e:
         print(e)
 
+
+
+# User Invoice Numbers
+
+class UserInvoiceNumberUpdate(UserPassesTestMixin, UpdateView):
+    template_name = "dispatch/userinvoicenumber/edit.html"
+    model = UserInvoiceNumber 
+    fields = ['number']
+
+    def get(self,request, *args, **kwargs):
+        if request.user.is_superuser:
+            self.fields.insert(1,'user')
+        return super(UserInvoiceNumberUpdate, self).get(request)
+
+    def test_func(self):
+        return self.request.user.is_superuser or \
+            self.get_object().user.pk is self.request.user.pk
+
+
 # Identity
 class IdentityDetail(DetailView):
     template_name = "dispatch/identity/detail.html"
     model = Identity
 
-def SetDefaultIdentity(request, pk):
-    ident = Ident.objects.get(pk=pk)
-    ident.set_default()
-    return redirect(reverse('identity_detail', kwargs={'pk': pk}))
+def SetDefaultIdentity(request, user_id, pk):
+    if request.user.is_superuser:
+        print("WORKS")
+        ident = Identity.objects.get(pk=pk)
+        ident.set_default()
+    return redirect(reverse('driver_details', kwargs={'pk': user_id}))
 
 class IdentityCreate(UserPassesTestMixin, CreateView):
     template_name = "dispatch/identity/edit.html"