QA find basket: add treatment file

3 files changed, 30 insertions, 0 deletions

diff --git a/archaeological_finds/models_finds.py b/archaeological_finds/models_finds.py
index 24e5227c1..2616d61fa 100644
--- a/archaeological_finds/models_finds.py
+++ b/archaeological_finds/models_finds.py
@@ -674,6 +674,11 @@ class FindBasket(Basket, MainItem, ValueGetter):
                 (reverse('findbasket-add-treatment', args=[self.pk]),
                  _(u"Add treatment"), "fa fa-flask", "", "", False),
             ]
+        if self.can_do(request, "add_treatmentfile"):
+            actions += [
+                (reverse('findbasket-add-treatmentfile', args=[self.pk]),
+                 _(u"Add treatment file"), "fa fa-file-text-o", "", "", False),
+            ]
 
         duplicate = self.QUICK_ACTIONS[0]
         actions += [
diff --git a/archaeological_finds/urls.py b/archaeological_finds/urls.py
index 95a20e3ff..7b70f2296 100644
--- a/archaeological_finds/urls.py
+++ b/archaeological_finds/urls.py
@@ -93,6 +93,11 @@ urlpatterns = [
         check_rights(['change_find', 'change_own_find'])(
             views.findbasket_treatment_add),
         name='findbasket-add-treatment'),
+    url(r'^findbasket-add-treatmentfile/(?P<pk>[0-9-]+)/$',
+        check_rights(["add_treatmentfile", "add_own_treatmentfile"])(
+            views.findbasket_treatmentfile_add),
+        name='findbasket-add-treatmentfile'),
+
     url(r'^find-add-treatment/(?P<pk>[0-9-]+)/$',
         check_rights(['change_find', 'change_own_find'])(
             views.find_treatment_add),
diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index f0dfe28e7..8d46866be 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -650,6 +650,14 @@ def findbasket_treatment_add(request, pk, current_right=None):
         request, ",".join([str(f.pk) for f in basket.items.all()]))
 
 
+def findbasket_treatmentfile_add(request, pk, current_right=None):
+    try:
+        basket = models.FindBasket.objects.get(pk=pk)
+    except models.FindBasket.DoesNotExist:
+        raise Http404()
+    return treatmentfile_add(request, basket_pk=basket.pk)
+
+
 def container_treatment_add(request, pk, current_right=None):
     try:
         basket = models.FindBasket.objects.get(pk=pk)
@@ -780,6 +788,18 @@ def treatmentfile_modify(request, pk):
         'treatmentfile_modification',
         kwargs={'step': 'treatmentfile-treatmentfile_modification'}))
 
+
+def treatmentfile_add(request, basket_pk=None):
+    treatmentfile_creation_wizard(request)
+    wizards.TreatmentFileWizard.session_set_value(
+        request, 'treatmentfile-treatmentfile_creation',
+        'associated_basket', basket_pk, reset=True)
+    url = reverse('treatmentfile_creation',
+                  kwargs={'step': 'treatmentfile-treatmentfile_creation'})
+    url += "?associated_basket={}".format(basket_pk)
+    return redirect(url)
+
+
 treatmentfile_deletion_wizard = wizards.TreatmentFileDeletionWizard.as_view([
     ('selec-treatmentfile_deletion', forms.TreatmentFileFormSelection),
     ('final-treatmentfile_deletion', forms.TreatmentFileDeletionForm)],