TokenBatch.spend(), spendable(), unlocked_fraction

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):

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)
+        print(token_batches)
+        # print(o.deposit(100))
+        # print(o.token_price())
+        # print(o._collateral_pool)