From dd984028757eccdcceb8016d9eed04681dacc8e9 Mon Sep 17 00:00:00 2001
From: Andrew Chiw <randomshinichi4869@gmail.com>
Date: Mon, 27 Apr 2020 19:34:05 +0200
Subject: [PATCH] TokenBatch.spend(), spendable(), unlocked_fraction

---
 hatch.py      | 28 ++++++++++++++++++++++++++--
 hatch_test.py | 37 +++++++++++++++++++++++++++++++++----
 2 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/hatch.py b/hatch.py
index ab5253b..fd6ec0b 100644
--- a/hatch.py
+++ b/hatch.py
@@ -49,17 +49,41 @@ class TokenBatch:
         self.creation_date = datetime.today()
         self.cliff_days = cliff_days
         self.halflife_days = halflife_days
+        self.spent = 0
+
     def __repr__(self):
-        o = "TokenBatch {} {}, Unlocked: {}".format("Hatch" if self.hatch_tokens else "", self.value, self.unlocked(datetime.today()))
+        o = "TokenBatch {} {}, Unlocked: {}".format("Hatch" if self.hatch_tokens else "", self.value, self.unlocked_fraction(datetime.today()))
         return o
 
-    def unlocked(self, day: datetime = datetime.today()) -> float:
+    def unlocked_fraction(self, day: datetime = datetime.today()) -> float:
+        """
+        returns what fraction of the TokenBatch is unlocked to date
+        """
         if self.hatch_tokens:
             days_delta = day - self.creation_date
             u = vesting_curve(days_delta.days, self.cliff_days, self.halflife_days)
             return u if u > 0 else 0
         else:
             return 1.0
+    
+    def spend(self, x: float):
+        """
+        checks if you can spend so many tokens, then decreases this TokenBatch instance's value accordingly
+        returns the argument if successful for your convenience
+        """
+        if x > self.spendable():
+            raise Exception("Not so many tokens are available for you to spend yet!")
+
+        self.value -= x
+        self.spent += x
+        return x
+
+    def spendable(self) -> float:
+        """
+        spendable() = self.unlocked_fraction * self.value - self.spent
+        Needed in case some Tokens were burnt before.
+        """
+        return (self.unlocked_fraction() * self.value) - self.spent
 
 class Commons:
     def __init__(self, total_hatch_raise, token_supply, hatch_tribute=0.2, exit_tribute=0):
diff --git a/hatch_test.py b/hatch_test.py
index e249d0b..944d173 100644
--- a/hatch_test.py
+++ b/hatch_test.py
@@ -1,6 +1,6 @@
 from hatch import *
 import unittest
-from datetime import datetime, timedelta
+import datetime
 
 class TestHatch(unittest.TestCase):
     def test_vesting_curve(self):
@@ -11,6 +11,34 @@ class TestHatch(unittest.TestCase):
     def test_convert_80p_to_halflife(self):
         self.assertEqual(convert_80p_to_cliff_and_halflife(90), (41.64807836875666, 20.82403918437833))
 
+class TokenBatchTest(unittest.TestCase):
+    def test_unlocked_fraction(self):
+        tbh = TokenBatch(10000, 3, 3, True)
+        tb = TokenBatch(10000, 5, 10, False)
+
+        self.assertEqual(tbh.unlocked_fraction(), 0)
+        self.assertEqual(tbh.unlocked_fraction(datetime.datetime.today() + datetime.timedelta(days=3)), 0)
+        self.assertEqual(tbh.unlocked_fraction(datetime.datetime.today() + datetime.timedelta(days=6)), 0.5)
+
+        self.assertEqual(tb.unlocked_fraction(), 1.0)
+    
+    def test_spend(self):
+        tbh = TokenBatch(10000, 3, 3, True)
+
+        with self.assertRaises(Exception):
+            tbh.spend(100)
+
+        tb = TokenBatch(10000, 3, 3, False)
+        tb.spend(100)
+        self.assertEqual(tb.value, 9900)
+        self.assertEqual(tb.spent, 100)
+        tb.spend(1000)
+        self.assertEqual(tb.value, 8900)
+        self.assertEqual(tb.spent, 1100)
+
+        with self.assertRaises(Exception):
+            tb.spend(10000)
+
 class TestSystem(unittest.TestCase):
     def test_system(self):
         # 100,000 DAI invested for 1,000,000 tokens.
@@ -26,6 +54,7 @@ class TestSystem(unittest.TestCase):
         self.assertEqual(o._token_supply, token_supply_initial)
 
         self.assertEqual(o.token_price(), 0.14)
-        print(o.deposit(100))
-        print(o.token_price())
-        print(o._collateral_pool)
\ No newline at end of file
+        print(token_batches)
+        # print(o.deposit(100))
+        # print(o.token_price())
+        # print(o._collateral_pool)
\ No newline at end of file