Treatment: workon on basket update after treatment

1 files changed, 4 insertions, 10 deletions

diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 5f2b4346d..409f7b28f 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -33,22 +33,16 @@ class PackagingWizard(TreatmentWizard):
         kwargs['user'] = self.request.user
         return kwargs
 
-
     def save_model(self, dct, m2m, whole_associated_models, form_list,
                    return_object):
         dct = self.get_extra_model(dct, form_list)
         obj = self.get_current_saved_object()
         dct['location'] = dct['container'].location
-        basket = dct.pop('basket')
+        items = dct.pop('basket')
         treatment = Treatment(**dct)
-        treatment.save()
-        for item in basket.items.all():
-            new = item.duplicate(self.request.user)
-            item.downstream_treatment = treatment
-            item.save()
-            new.upstream_treatment = treatment
-            new.container = dct['container']
-            new.save()
+        extra_args_for_new = {"container": dct['container']}
+        treatment.save(items=items, user=self.request.user,
+                       extra_args_for_new=extra_args_for_new)
         res = render_to_response('ishtar/wizard/wizard_done.html', {},
                                  context_instance=RequestContext(self.request))
         return return_object and (obj, res) or res