QA find basket: add treatment file

1 files changed, 20 insertions, 0 deletions

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)],