support rewinding specific changes

2 files changed, 51 insertions, 3 deletions

diff --git a/geldschieberbot.py b/geldschieberbot.py
index d74cca5..7ae6454 100644
--- a/geldschieberbot.py
+++ b/geldschieberbot.py
@@ -190,7 +190,7 @@ cars pay car-name amount - pay a bill for the specified car
 
 tanken amount [person] [car] [info] - calculate fuel costs, service charge and add them to the person's and car's balance respectively
 
-fuck - rewind last change
+fuck [change] - rewind last or specific change
 list-changes [n] [first] - list the last n changes starting at first
 
 weekly name cmd -  repeat cmd each week
@@ -595,11 +595,22 @@ def fuck(sender, args, msg):
     else:
         name = num2name[sender]
 
-    if len(changes[name]) == 0:
+    nchanges = len(changes[name])
+    if nchanges == 0:
         return "Nothing to rewind", None
 
+    change_to_rewind = -1
+    if len(args) >= 2:
+        try:
+            change_to_rewind = int(args[1]) - 1
+        except ValueError:
+            return None, "change to rewind must be a number"
+
+        if change_to_rewind > nchanges:
+            return None, "change to rewind is bigger than there are changes"
+
     # pop last item
-    last_changes = changes[name].pop()
+    last_changes = changes[name].pop(change_to_rewind)
     args, last_changes = last_changes[0], last_changes[1:]
 
     output = name + ": sorry I fucked up!\nRewinding:\n"
diff --git a/test.py b/test.py
index 4d05bd6..17194fb 100755
--- a/test.py
+++ b/test.py
@@ -800,6 +800,16 @@ class TestFuckCmd(unittest.TestCase):
         res = run_bot(self, num[alice], "!fuck")
         self.assertEqual(res.stdout, 'Nothing to rewind')
 
+    def test_fuck_invalid_change(self):
+        run_bot(self, num[alice], f"!schieb {bob} 10")
+        res = run_bot(self, num[alice], f"!fuck foo")
+        msg = "ERROR: change to rewind must be a number"
+        self.assertEqual(res.stdout, msg)
+
+        res = run_bot(self, num[alice], f"!fuck 20")
+        msg = "ERROR: change to rewind is bigger than there are changes"
+        self.assertEqual(res.stdout, msg)
+
     def test_fuck_transaction(self):
         for cmd in ["fuck", "undo", "rewind"]:
             run_bot(self, num[alice], f"!schieb {bob} 10")
@@ -813,6 +823,33 @@ alice <- 10.00 bob
             self.assertEqual(res.stdout, msg)
             compare_state("test/state.json_3users")
 
+    def test_rewind_first(self):
+        run_bot(self, num[alice], f"!split 3 {bob} {charlie}")
+        run_bot(self, num[alice], f"!schieb {bob} 10")
+        run_bot(self, num[alice], f"!nimm {charlie} 10")
+        res = run_bot(self, num[alice], f"!fuck 1")
+        msg = ""\
+"""alice: sorry I fucked up!
+Rewinding:
+!split 3 bob charlie
+alice <- 1.00 bob
+alice <- 1.00 charlie
+"""
+        self.assertEqual(res.stdout, msg)
+
+    def test_rewind_second(self):
+        run_bot(self, num[alice], f"!split 10 {bob} {charlie}")
+        run_bot(self, num[alice], f"!schieb {bob} 10")
+        run_bot(self, num[alice], f"!nimm {charlie} 10")
+        res = run_bot(self, num[alice], f"!fuck 2")
+        msg = ""\
+"""alice: sorry I fucked up!
+Rewinding:
+!schieb bob 10
+alice <- 10.00 bob
+"""
+        self.assertEqual(res.stdout, msg)
+
     def test_fuck_split(self):
         run_bot(self, num[alice], "!split 3 bob charlie")
         res = run_bot(self, num[alice], "!fuck")