Upgrade logic for GPUCreditTracker class

Author: nwthomas

Makefile

@@ -18,4 +18,4 @@ test-js:
 
 test-py:
 	@echo "Running all Python tests..."
-	uv run pytest
+	uv run pytest -s

src/miscellaneous-code-challenges/gpu-credit-tracker/gpu_credit_tracker.py

@@ -3,6 +3,13 @@
 from heapq import heappop, heappush
 from typing import List, Optional
 
+from enum import Enum
+
+
+class Action(Enum):
+    GRANT = "grant"
+    SUBTRACT = "subtract"
+
 
 @dataclass
 class Transaction:
@@ -11,72 +18,64 @@ class Transaction:
     at a given point in time
     """
 
-    action: str
-    amount: int
+    action: Action
     timestamp: int
-    expiration_timestamp: Optional[int]
-
-    def __init__(
-        self,
-        action: str,
-        timestamp: int,
-        amount: int,
-        expiration_timestamp: Optional[int],
-    ):
-        self.action = action
-        self.amount = amount
-        self.timestamp = timestamp
-        self.expiration_timestamp = expiration_timestamp
+    amount: int
+    expiration_timestamp: Optional[int] = field(default=None)
 
     def __lt__(self, other_value) -> bool:
+        """Compares two transactions. Required for heap comparisons. Only works with grants."""
+        if self.action != Action.GRANT:
+            return False
+
         return self.expiration_timestamp < other_value.expiration_timestamp
 
 
-@dataclass
 class GPUCreditTracker:
     """
     Implements the ability to create GPU credit grants, subtract credits, and get the credit
     balance at a given point in time.
     """
 
-    transactions: List[Transaction] = field(default_factory=list)
+    def __init__(self):
+        self.transactions: List[Transaction] = []
 
     def create_grant(
         self, timestamp: int, expiration_timestamp: int, amount: int
     ) -> None:
-        """Creates a grant for a credit balance increase"""
-        transaction = Transaction("grant", timestamp, amount, expiration_timestamp)
+        """Creates a grant for a credit balance increase."""
+        transaction = Transaction(Action.GRANT, timestamp, amount, expiration_timestamp)
         self.__insert_transaction(transaction)
 
     def subtract(self, timestamp: int, amount: int) -> None:
-        """Subtracts credits from the current balance if possible"""
-        transaction = Transaction("subtract", timestamp, amount, None)
+        """Subtracts credits from the current balance if possible."""
+        transaction = Transaction(Action.SUBTRACT, timestamp, amount, None)
         self.__insert_transaction(transaction)
 
-    def get_balance(self, timestamp: int) -> int:
-        """Returns the balance at a certain point in time"""
-        grant_heap = []
-        current = 0
+    def get_balance(self, timestamp: int) -> int | None:
+        """Returns the balance at a certain point in time."""
+        has_gone_negative = False
+        grant_heap: list[Transaction] = []
+        i = 0
 
         while (
-            current < len(self.transactions)
-            and self.transactions[current].timestamp <= timestamp
+            i < len(self.transactions) and self.transactions[i].timestamp <= timestamp
         ):
-            current_transaction = deepcopy(self.transactions[current])
+            current_transaction = deepcopy(self.transactions[i])
 
             # Clear out stale grants that have expired and can no longer be used
             while (
-                len(grant_heap) > 0
+                grant_heap
                 and grant_heap[0].expiration_timestamp < current_transaction.timestamp
             ):
                 heappop(grant_heap)
 
-            if current_transaction.action == "grant":
+            if current_transaction.action == Action.GRANT:
                 heappush(grant_heap, current_transaction)
-            elif current_transaction.action == "subtract":
+            elif current_transaction.action == Action.SUBTRACT:
                 total_subtract_remaining = current_transaction.amount
 
-                while len(grant_heap) > 0 and total_subtract_remaining > 0:
+                while grant_heap and total_subtract_remaining > 0:
                     current_grant_remaining_amount = grant_heap[0].amount
 
                     if current_grant_remaining_amount >= total_subtract_remaining:
@@ -86,39 +85,49 @@ def get_balance(self, timestamp: int) -> int:
                         total_subtract_remaining -= grant_heap[0].amount
                         heappop(grant_heap)
 
-            current += 1
+                # Any leftover balance at a timestamp means that the account is in a not recoverable
+                # state and None should always be returned from this point forward. Exit immediately.
+                if total_subtract_remaining > 0:
+                    has_gone_negative = True
+                    break
+
+            i += 1
 
+        # Calculate the remaining balance from the grant heap if not in unrecoverable state
         total = 0
-        while len(grant_heap) > 0:
+        while not has_gone_negative and grant_heap:
             current_grant = heappop(grant_heap)
 
             if current_grant.expiration_timestamp > timestamp:
                 total += current_grant.amount
 
-        return total if total > 0 else 0
+        return total if not has_gone_negative else None
 
     def __insert_transaction(self, transaction: Transaction) -> None:
-        if len(self.transactions) == 0:
-            self.transactions.append(transaction)
-        else:
-            current = 0
-
-            while current < len(self.transactions) and (
-                (
-                    # Place grant transactions at beginning of same timestamp
-                    transaction.action == "grant"
-                    and self.transactions[current].timestamp < transaction.timestamp
-                )
-                or (
-                    # Place subtract transactions at end of same timestamp
-                    transaction.action == "subtract"
-                    and self.transactions[current].timestamp <= transaction.timestamp
+        """Inserts a new transaction into the list of transactions."""
+        new_transaction: Transaction | None = transaction
+        new_transactions_list: list[Transaction] = []
+
+        for i, t in enumerate(self.transactions):
+            if i == 0 and new_transaction and new_transaction.timestamp < t.timestamp:
+                new_transactions_list.insert(0, new_transaction)
+                new_transaction = None
+
+            new_transactions_list.append(t)
+
+            if (
+                i + 1 < len(self.transactions)
+                and new_transaction
+                and (
+                    t.timestamp
+                    < new_transaction.timestamp
+                    < self.transactions[i + 1].timestamp
                 )
             ):
-                current += 1
+                new_transactions_list.append(new_transaction)
+                new_transaction = None
+
+        if new_transaction:
+            new_transactions_list.append(new_transaction)
 
-            self.transactions = (
-                self.transactions[:current]
-                + [transaction]
-                + self.transactions[current:]
-            )
+        self.transactions = new_transactions_list

src/miscellaneous-code-challenges/gpu-credit-tracker/test_gpu_credit_tracker.py

@@ -1,45 +1,94 @@
-import unittest
-
 from gpu_credit_tracker import GPUCreditTracker
 
 
-class TestGPUCreditTracker(unittest.TestCase):
-    def test_inserts_correctly(self):
-        """Tests if the class inserts transactions correctly and gets the balance"""
-        tracker = GPUCreditTracker()
-        tracker.create_grant(10, 50, 5)
-        tracker.create_grant(20, 60, 5)
-        self.assertEqual(tracker.get_balance(19), 5)
-        self.assertEqual(tracker.get_balance(20), 10)
-        self.assertEqual(tracker.get_balance(50), 5)
-        self.assertEqual(tracker.get_balance(60), 0)
-
-    def test_subtracts_in_any_order_entered(self):
-        """Tests if subtraction can be handled in any order that it's received"""
-        tracker = GPUCreditTracker()
-        tracker.create_grant(10, 50, 5)
-        tracker.create_grant(20, 60, 5)
-        self.assertEqual(tracker.get_balance(19), 5)
-        tracker.subtract(15, 4)
-        self.assertEqual(tracker.get_balance(10), 5)
-        self.assertEqual(tracker.get_balance(20), 6)
-        tracker.create_grant(15, 20, 8)
-        self.assertEqual(tracker.get_balance(14), 5)
-        self.assertEqual(tracker.get_balance(15), 9)
-        self.assertEqual(tracker.get_balance(20), 10)
-
-    def test_should_return_zero_for_no_grants(self):
-        """Returns 0 for get_balance() if no grants have been created"""
-        tracker = GPUCreditTracker()
-        self.assertEqual(tracker.get_balance(10), 0)
-
-    def test_handles_all_grants_expired(self):
-        """Returns 0 if all grants have been expired"""
-        tracker = GPUCreditTracker()
-        tracker.create_grant(10, 50, 5)
-        tracker.create_grant(20, 60, 5)
-        self.assertEqual(tracker.get_balance(100), 0)
-
-
-if __name__ == "__main__":
-    unittest.main()
+def test_inserts_correctly():
+    """Tests if the class inserts transactions correctly and gets the balance"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    credit.create_grant(20, 60, 5)
+    assert credit.get_balance(19) == 5
+    assert credit.get_balance(20) == 10
+    assert credit.get_balance(50) == 5
+    assert credit.get_balance(60) == 0
+
+
+def test_subtracts_from_expiring_grants_first():
+    """Tests if subtraction can be handled from expiring grants first"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    credit.create_grant(14, 60, 10)
+    assert credit.get_balance(15) == 15
+    credit.subtract(15, 10)
+    assert credit.get_balance(10) == 5
+    assert credit.get_balance(20) == 5
+
+
+def test_subtracts_in_any_order_entered():
+    """Tests if subtraction can be handled in any order that it's received"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    credit.create_grant(20, 60, 5)
+    assert credit.get_balance(19) == 5
+    credit.subtract(15, 4)
+    assert credit.get_balance(10) == 5
+    assert credit.get_balance(20) == 6
+    credit.create_grant(14, 20, 8)
+    assert credit.get_balance(14) == 13
+    assert credit.get_balance(15) == 9
+    assert credit.get_balance(20) == 10
+
+
+def test_should_return_zero_for_no_grants():
+    """Returns 0 for get_balance() if no grants have been created"""
+    credit = GPUCreditTracker()
+    assert credit.get_balance(10) == 0
+
+
+def test_handles_all_grants_expired():
+    """Returns 0 if all grants have been expired"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    credit.create_grant(20, 60, 5)
+    assert credit.get_balance(100) == 0
+
+
+def test_handles_negative_balance():
+    """Returns None if the balance goes negative"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    credit.subtract(15, 6)
+    credit.create_grant(16, 20, 8)
+    assert credit.get_balance(15) == None
+    assert credit.get_balance(16) == None
+    assert credit.get_balance(20) == None
+    credit.create_grant(11, 100, 10)
+    assert credit.get_balance(10) == 5
+    assert credit.get_balance(11) == 15
+    assert credit.get_balance(15) == 9
+    assert credit.get_balance(16) == 17
+    assert credit.get_balance(20) == 9
+
+
+def test_balance_before_first_transaction():
+    """Returns 0 for get_balance(ts) when ts is before the first transaction"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    assert credit.get_balance(5) == 0
+    assert credit.get_balance(9) == 0
+
+
+def test_exact_expiration_boundary():
+    """Grant is not counted at the expiration instant (expiration is exclusive)"""
+    credit = GPUCreditTracker()
+    credit.create_grant(10, 50, 5)
+    assert credit.get_balance(49) == 5
+    assert credit.get_balance(50) == 0
+
+
+def test_insert_transaction_places_earlier_transaction_at_start():
+    """Transaction with timestamp before all existing ones is inserted at start, not end."""
+    credit = GPUCreditTracker()
+    credit.create_grant(20, 60, 5)
+    credit.create_grant(10, 50, 3)
+    assert credit.get_balance(15) == 3
+    assert credit.get_balance(25) == 8