Lots of work on the stateful checking stuff: it's a fair bit nicer now

1 files changed, 20 insertions, 2 deletions

diff --git a/src/main/java/au/id/zancanaro/javacheck/state/CommandValue.java b/src/main/java/au/id/zancanaro/javacheck/state/CommandValue.java
index 8d3f272..8584fc7 100644
--- a/src/main/java/au/id/zancanaro/javacheck/state/CommandValue.java
+++ b/src/main/java/au/id/zancanaro/javacheck/state/CommandValue.java
@@ -5,13 +5,31 @@ import java.util.NoSuchElementException;
 import java.util.function.Supplier;
 
 public class CommandValue<T> {
+    public static interface Action<T> {
+        T doAction() throws Throwable;
+    }
+
+    public static interface VoidAction {
+        void doAction() throws Throwable;
+    }
+
     private static Map<Integer, Object> values = null;
 
-    public static <T> T withValues(Map<Integer, Object> newValues, Supplier<T> action) {
+    public static <T> T withValues(Map<Integer, Object> newValues, Action<T> action) throws Throwable {
+        Map<Integer,Object> oldValues = values;
+        try {
+            values = newValues;
+            return action.doAction();
+        } finally {
+            values = oldValues;
+        }
+    }
+
+    public static void withValues(Map<Integer, Object> newValues, VoidAction action) throws Throwable {
         Map<Integer,Object> oldValues = values;
         try {
             values = newValues;
-            return action.get();
+            action.doAction();
         } finally {
             values = oldValues;
         }