1 files changed, 26 insertions, 6 deletions
diff --git a/src/test/java/org/zwobble/mammoth/tests/Arguments.java b/src/test/java/org/zwobble/mammoth/tests/Arguments.java
index afb20bb..badb35a 100644
--- a/src/test/java/org/zwobble/mammoth/tests/Arguments.java
+++ b/src/test/java/org/zwobble/mammoth/tests/Arguments.java
@@ -18,18 +18,38 @@ public class Arguments {
     }
 
     public <T> T get(Class<T> clazz, T defaultValue) {
-        Iterator<T> arguments = getAll(clazz).iterator();
-        if (!arguments.hasNext()) {
-            return defaultValue;
-        } else {
-            return only(arguments);
-        }
+        return onlyOrDefault(getAll(clazz), defaultValue);
+    }
+
+    public <T> T get(ArgumentKey<T> key, T defaultValue) {
+        return onlyOrDefault(getAll(key), defaultValue);
     }
 
     private <T> Iterable<T> getAll(Class<T> clazz) {
         return Iterables.lazyFilter(asList(arguments), clazz);
     }
 
+    private <T> Iterable<T> getAll(ArgumentKey<T> key) {
+        return Iterables.lazyMap(
+            Iterables.lazyFilter(
+                Iterables.lazyFilter(asList(arguments), Argument.class),
+                argument -> argument.getKey().equals(key)
+            ),
+            argument -> (T)argument.getValue()
+        );
+    }
+
+    private <T> T onlyOrDefault(Iterable<T> iterable, T defaultValue) {
+        return onlyOrDefault(iterable.iterator(), defaultValue);
+    }
+
+    private <T> T onlyOrDefault(Iterator<T> iterator, T defaultValue) {
+        if (!iterator.hasNext()) {
+            return defaultValue;
+        } else {
+            return only(iterator);
+        }
+    }
 
     private <T> T only(Iterable<T> iterable) {
         return only(iterable.iterator());