From eb615ab88551bfeca14266d7081cfc72dac419c0 Mon Sep 17 00:00:00 2001
From: Jeferson Moura <jefmmoreira@gmail.com>
Date: Fri, 17 Aug 2018 10:14:04 +0200
Subject: [PATCH] Fix Budget signal wasn't calculating correctly

---
 workflow/signals.py            | 20 ++++++++++----------
 workflow/tests/test_signals.py | 26 +++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/workflow/signals.py b/workflow/signals.py
index 12629631b..97819ee09 100644
--- a/workflow/signals.py
+++ b/workflow/signals.py
@@ -330,22 +330,22 @@ def sync_delete_track_workflowlevel1(sender, instance, **kwargs):
 
 
 # BUDGET SIGNALS
-@receiver(signals.post_save, sender=Budget)
+@receiver(signals.pre_save, sender=Budget)
 def save_budget_value_from_wfl2(sender, instance, *args, **kwargs):
     if instance.workflowlevel2 is not None:
         wflvl2 = instance.workflowlevel2
-
-        if kwargs.get('created'):
-            # Sum the new values
-            wflvl2.total_estimated_budget += instance.proposed_value
-            wflvl2.actual_cost += instance.actual_value
-        else:
-            # Subtract the old values
+        try:
             old_budget = Budget.objects.get(id=instance.id)
+            # Subtract the old values
             wflvl2.total_estimated_budget -= old_budget.proposed_value
             wflvl2.actual_cost -= old_budget.actual_value
-
-        wflvl2.save()
+        except Budget.DoesNotExist:
+            pass
+        finally:
+            # Sum the new values
+            wflvl2.total_estimated_budget += instance.proposed_value
+            wflvl2.actual_cost += instance.actual_value
+            wflvl2.save()
 
 
 @receiver(signals.post_delete, sender=Budget)
diff --git a/workflow/tests/test_signals.py b/workflow/tests/test_signals.py
index dd837d332..bf07156e3 100644
--- a/workflow/tests/test_signals.py
+++ b/workflow/tests/test_signals.py
@@ -987,7 +987,7 @@ def test_save_budget_value_from_wfl2_create(self):
         self.assertEqual(wflvl2.total_estimated_budget, 100.00)
         self.assertEqual(wflvl2.actual_cost, 10.00)
 
-    def test_save_budget_value_from_wfl2_update(self):
+    def test_save_budget_value_from_wfl2_create_more(self):
         """
         When a budget is updated, the related wfl2s total_estimated_budget
         should be updated
@@ -1011,6 +1011,30 @@ def test_save_budget_value_from_wfl2_update(self):
         self.assertEqual(wflvl2.total_estimated_budget, 180.00)
         self.assertEqual(wflvl2.actual_cost, 30.00)
 
+    def test_save_budget_value_from_wfl2_update(self):
+        """
+        When a budget is updated, the related wfl2s total_estimated_budget
+        should be updated
+        """
+
+        wflvl1 = factories.WorkflowLevel1(name='WorkflowLevel1',
+                                          organization=self.organization)
+        wflvl2 = factories.WorkflowLevel2(name='WorkflowLevel2',
+                                          workflowlevel1=wflvl1,
+                                          total_estimated_budget=0,
+                                          actual_cost=0)
+        budget = factories.Budget(proposed_value=100,
+                                  actual_value=10,
+                                  workflowlevel2=wflvl2)
+
+        budget.proposed_value = 150
+        budget.actual_value = 50
+        budget.save()
+
+        wflvl2 = WorkflowLevel2.objects.get(pk=wflvl2.pk)
+        self.assertEqual(wflvl2.total_estimated_budget, 150.00)
+        self.assertEqual(wflvl2.actual_cost, 50.00)
+
 
 @tag('pkg')
 class BudgetDeleteTest(TestCase):