From 2b071a1a3bc752bfaa0fd6c0b2d8f8460721d6d8 Mon Sep 17 00:00:00 2001
From: Jocelyn Delande <jocelyn@crapouillou.net>
Date: Thu, 20 Aug 2015 10:33:43 +0200
Subject: [PATCH] Add members weight in models and budget backend refs #94

---
 budget/forms.py                             |  2 ++
 budget/migrations/versions/26d6a218c329_.py | 26 ++++++++++++++++++
 budget/models.py                            |  9 ++++---
 budget/tests.py                             | 30 +++++++++++++++++++++
 4 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 budget/migrations/versions/26d6a218c329_.py

diff --git a/budget/forms.py b/budget/forms.py
index 2dde57d..918e82a 100644
--- a/budget/forms.py
+++ b/budget/forms.py
@@ -152,6 +152,7 @@ class BillForm(Form):
 class MemberForm(Form):
 
     name = TextField(_("Name"), validators=[Required()])
+    weight = CommaDecimalField(_("Weight"), default=1)
     submit = SubmitField(_("Add"))
 
     def __init__(self, project, *args, **kwargs):
@@ -170,6 +171,7 @@ class MemberForm(Form):
         # if the user is already bound to the project, just reactivate him
         person.name = self.name.data
         person.project = project
+        person.weight = self.weight.data
 
         return person
 
diff --git a/budget/migrations/versions/26d6a218c329_.py b/budget/migrations/versions/26d6a218c329_.py
new file mode 100644
index 0000000..859b9af
--- /dev/null
+++ b/budget/migrations/versions/26d6a218c329_.py
@@ -0,0 +1,26 @@
+"""Add Person.weight column
+
+Revision ID: 26d6a218c329
+Revises: b9a10d5d63ce
+Create Date: 2016-06-15 09:22:04.069447
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = '26d6a218c329'
+down_revision = 'b9a10d5d63ce'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('person', sa.Column('weight', sa.Float(), nullable=True))
+    ### end Alembic commands ###
+
+
+def downgrade():
+    ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('person', 'weight')
+    ### end Alembic commands ###
diff --git a/budget/models.py b/budget/models.py
index 727200f..16cc6c1 100644
--- a/budget/models.py
+++ b/budget/models.py
@@ -40,8 +40,9 @@ class Project(db.Model):
             bills = Bill.query.filter(Bill.owers.contains(person))
             for bill in bills.all():
                 if person != bill.payer:
-                    should_pay[person] += bill.pay_each()
-                    should_receive[bill.payer] += bill.pay_each()
+                    share = bill.pay_each() * person.weight
+                    should_pay[person] += share
+                    should_receive[bill.payer] += share
 
         for person in self.members:
             balance = should_receive[person] - should_pay[person]
@@ -159,6 +160,7 @@ class Person(db.Model):
     bills = db.relationship("Bill", backref="payer")
 
     name = db.Column(db.UnicodeText)
+    weight = db.Column(db.Float, default=1)
     activated = db.Column(db.Boolean, default=True)
 
     def has_bills(self):
@@ -219,7 +221,8 @@ class Bill(db.Model):
     def pay_each(self):
         """Compute what each person has to pay"""
 	if self.owers:
-		return self.amount / len(self.owers)
+                # FIXME: SQL might dot that more efficiently
+		return self.amount / sum(i.weight for i in self.owers)
 	else:
 		return 0
 
diff --git a/budget/tests.py b/budget/tests.py
index 760ffc0..021b425 100644
--- a/budget/tests.py
+++ b/budget/tests.py
@@ -416,6 +416,36 @@ class BudgetTestCase(TestCase):
         bill = models.Bill.query.filter(models.Bill.date == '2011-08-01')[0]
         self.assertEqual(bill.amount, 25.02)
 
+    def test_weighted_balance(self):
+        self.post_project("raclette")
+
+        # add two persons
+        self.app.post("/raclette/members/add", data={'name': 'alexis'})
+        self.app.post("/raclette/members/add", data={'name': 'freddy familly', 'weight': 4})
+
+        members_ids = [m.id for m in
+                       models.Project.query.get("raclette").members]
+
+        # test balance
+        self.app.post("/raclette/add", data={
+            'date': '2011-08-10',
+            'what': u'fromage à raclette',
+            'payer': members_ids[0],
+            'payed_for': members_ids,
+            'amount': '10',
+        })
+
+        self.app.post("/raclette/add", data={
+            'date': '2011-08-10',
+            'what': u'pommes de terre',
+            'payer': members_ids[1],
+            'payed_for': members_ids,
+            'amount': '10',
+        })
+
+        balance = models.Project.query.get("raclette").balance
+        self.assertEqual(set(balance.values()), set([6, -6]))
+
     def test_rounding(self):
         self.post_project("raclette")