From ca8b9028bc4db6ae2c97456bbf662d2196c77aee Mon Sep 17 00:00:00 2001
From: dandds <dan-ctr@friends.dds.mil>
Date: Tue, 18 Sep 2018 09:54:28 -0400
Subject: [PATCH] reports shows real budget information for total spend

---
 atst/domain/reports.py          | 21 +++++++++------------
 atst/models/task_order.py       | 16 ++++++++++++++++
 atst/routes/workspaces.py       |  2 +-
 tests/domain/test_reports.py    | 18 ++++++++++++++++++
 tests/models/test_task_order.py | 12 ++++++++++++
 5 files changed, 56 insertions(+), 13 deletions(-)
 create mode 100644 tests/domain/test_reports.py

diff --git a/atst/domain/reports.py b/atst/domain/reports.py
index 7e49e0de..36e0bd52 100644
--- a/atst/domain/reports.py
+++ b/atst/domain/reports.py
@@ -154,18 +154,15 @@ CUMULATIVE_BUDGET_BELUGA = {
 
 class Reports:
     @classmethod
-    def workspace_totals(cls, alternate):
-        data = MONTHLY_SPEND_BELUGA if alternate else MONTHLY_SPEND_AARDVARK
-        spent = sum(
-            [
-                spend
-                for project in data.values()
-                for env in project.values()
-                for spend in env.values()
-            ]
-        )
-        budget = 70_000 if alternate else 500_000
-        return {"budget": budget, "spent": spent}
+    def workspace_totals(cls, workspace):
+        if workspace.request and workspace.request.task_order:
+            ws_to = workspace.request.task_order
+            budget = ws_to.budget
+        else:
+            budget = 0
+
+        # spent will be derived from CSP data
+        return {"budget": budget, "spent": 0}
 
     @classmethod
     def monthly_totals(cls, alternate):
diff --git a/atst/models/task_order.py b/atst/models/task_order.py
index 349d607c..863a189f 100644
--- a/atst/models/task_order.py
+++ b/atst/models/task_order.py
@@ -47,3 +47,19 @@ class TaskOrder(Base):
             for c in self.__table__.columns
             if c.name not in ["id", "attachment_id"]
         }
+
+    @property
+    def budget(self):
+        return sum(
+            filter(
+                None,
+                [
+                    self.clin_0001,
+                    self.clin_0003,
+                    self.clin_1001,
+                    self.clin_1003,
+                    self.clin_2001,
+                    self.clin_2003,
+                ],
+            )
+        )
diff --git a/atst/routes/workspaces.py b/atst/routes/workspaces.py
index 673719a5..4ee51150 100644
--- a/atst/routes/workspaces.py
+++ b/atst/routes/workspaces.py
@@ -109,7 +109,7 @@ def workspace_reports(workspace_id):
     return render_template(
         "workspaces/reports/index.html",
         cumulative_budget=Reports.cumulative_budget(alternate_reports),
-        workspace_totals=Reports.workspace_totals(alternate_reports),
+        workspace_totals=Reports.workspace_totals(workspace),
         monthly_totals=Reports.monthly_totals(alternate_reports),
         current_month=current_month,
         prev_month=prev_month,
diff --git a/tests/domain/test_reports.py b/tests/domain/test_reports.py
new file mode 100644
index 00000000..dcb0bfa8
--- /dev/null
+++ b/tests/domain/test_reports.py
@@ -0,0 +1,18 @@
+from atst.domain.reports import Reports
+
+from tests.factories import RequestFactory, TaskOrderFactory, WorkspaceFactory
+
+CLIN_NUMS = ["0001", "0003", "1001", "1003", "2001", "2003"]
+
+
+def test_workspace_totals():
+    task_order = TaskOrderFactory.create()
+
+    for num in CLIN_NUMS:
+        setattr(task_order, "clin_{}".format(num), 200)
+
+    request = RequestFactory.create(task_order=task_order)
+    workspace = WorkspaceFactory.create(request=request)
+    report = Reports.workspace_totals(workspace)
+    total = 200 * len(CLIN_NUMS)
+    assert report == {"budget": total, "spent": 0}
diff --git a/tests/models/test_task_order.py b/tests/models/test_task_order.py
index 900446a0..dd83a0ca 100644
--- a/tests/models/test_task_order.py
+++ b/tests/models/test_task_order.py
@@ -7,3 +7,15 @@ def test_as_dictionary():
     data = TaskOrderFactory.dictionary()
     real_task_order = TaskOrderFactory.create(**data)
     assert real_task_order.to_dictionary() == data
+
+
+def test_budget():
+    task_order = TaskOrderFactory.create(
+        clin_0001=500,
+        clin_0003=200,
+        clin_1001=None,
+        clin_1003=None,
+        clin_2001=None,
+        clin_2003=None,
+    )
+    assert task_order.budget == 700