Better code, based off ParentRunner among other things

6 files changed, 191 insertions, 102 deletions

diff --git a/src/main/java/au/id/zancanaro/Generator.java b/src/main/java/au/id/zancanaro/Generator.java
index 4b80e51..1c948af 100644
--- a/src/main/java/au/id/zancanaro/Generator.java
+++ b/src/main/java/au/id/zancanaro/Generator.java
@@ -1,7 +1,5 @@
 package au.id.zancanaro;
 
-import java.util.ArrayList;
-import java.util.List;
 import java.util.Random;
 import java.util.function.Function;
 
@@ -12,18 +10,6 @@ public interface Generator<T> {
         return (random, size) -> RoseTree.pure(value);
     }
 
-    default <R> Generator<R> genFlatmap(Function<RoseTree<T>, Generator<R>> f) {
-        return (random, size) -> {
-            RoseTree<T> inner = this.generate(random, size);
-            Generator<R> generator = f.apply(inner);
-            return generator.generate(random, size);
-        };
-    }
-
-    default <R> Generator<R> genFmap(Function<RoseTree<T>, RoseTree<R>> f) {
-        return (random, size) -> f.apply(this.generate(random, size));
-    }
-
     @SafeVarargs
     static <T> Generator<T[]> tuple(Generator<T>... generators) {
         return (random, size) -> {
@@ -42,10 +28,8 @@ public interface Generator<T> {
     }
 
     default <R> Generator<R> flatMap(Function<T, Generator<R>> action) {
-        return this.genFlatmap(rose -> {
-            Generator<RoseTree<R>> generator = (random, size) ->
-                    rose.fmap(action).fmap(g -> g.generate(random, size));
-            return generator.<R>genFmap(RoseTree::join);
-        });
+        return (random, size) -> {
+            return RoseTree.join(this.generate(random, size).fmap(action).fmap(g -> g.generate(random, size)));
+        };
     }
 }
diff --git a/src/main/java/au/id/zancanaro/Generators.java b/src/main/java/au/id/zancanaro/Generators.java
index f065d1d..5b025fa 100644
--- a/src/main/java/au/id/zancanaro/Generators.java
+++ b/src/main/java/au/id/zancanaro/Generators.java
@@ -1,14 +1,21 @@
 package au.id.zancanaro;
 
 import java.util.Iterator;
+import java.util.function.Predicate;
 
 public class Generators {
-    @SafeVarargs
-    public static <T> Generator<T[]> arrayGenerator(Generator<? extends T>... generators) {
-        return Generator.tuple((Generator<T>[]) generators);
+    public static <T> Generator<T> suchThat(Generator<T> gen, Predicate<T> pred) {
+        return (random, size) -> {
+            RoseTree<T> result = gen.generate(random, size);
+            if (pred.test(result.getValue())) {
+                return result.filter(pred);
+            } else {
+                return suchThat(gen, pred).generate(random, size);
+            }
+        };
     }
 
-    public static Generator<Integer> integerGenerator() {
+    public static Generator<Integer> integer() {
         return (random, size) -> {
             int value = random.nextInt(size);
             return new RoseTree<>(value, intShrinkingIterable(value));
diff --git a/src/main/java/au/id/zancanaro/Iterators.java b/src/main/java/au/id/zancanaro/Iterators.java
index b879c76..08cc37e 100644
--- a/src/main/java/au/id/zancanaro/Iterators.java
+++ b/src/main/java/au/id/zancanaro/Iterators.java
@@ -1,7 +1,11 @@
 package au.id.zancanaro;
 
+import java.util.Collections;
 import java.util.Iterator;
+import java.util.List;
+import java.util.Optional;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 public final class Iterators {
     public static <T> RangeIterator<T> rangeIterator(int countTo, Function<Integer,T> fn) {
@@ -127,4 +131,45 @@ public final class Iterators {
         }
 
     }
+
+    public static <T> FilteringIterator<T> filteringIterator(Predicate<T> pred, Iterator<T> iterator) {
+        return new FilteringIterator<>(pred, iterator);
+    }
+
+    private static class FilteringIterator<T> implements Iterator<T> {
+        private final Predicate<T> predicate;
+        private final Iterator<T> iterator;
+        private List<T> nextValue;
+
+        public FilteringIterator(Predicate<T> pred, Iterator<T> iterator) {
+            this.predicate = pred;
+            this.iterator = iterator;
+            this.nextValue = null;
+        }
+
+        private void populateNext() {
+            while (nextValue == null && iterator.hasNext()) {
+                T value = iterator.next();
+                if (predicate.test(value)) {
+                    nextValue = Collections.singletonList(value);
+                } else {
+                    nextValue = null;
+                }
+            }
+        }
+
+        @Override
+        public boolean hasNext() {
+            populateNext();
+            return nextValue != null;
+        }
+
+        @Override
+        public T next() {
+            populateNext();
+            T result = nextValue.get(0);
+            nextValue = null;
+            return result;
+        }
+    }
 }
diff --git a/src/main/java/au/id/zancanaro/PropertyTestRunner.java b/src/main/java/au/id/zancanaro/PropertyTestRunner.java
index fc90b8e..8df9b2a 100644
--- a/src/main/java/au/id/zancanaro/PropertyTestRunner.java
+++ b/src/main/java/au/id/zancanaro/PropertyTestRunner.java
@@ -2,97 +2,115 @@ package au.id.zancanaro;
 
 import au.id.zancanaro.annotations.Property;
 import au.id.zancanaro.annotations.Seed;
-import junit.framework.AssertionFailedError;
 import org.junit.AssumptionViolatedException;
+import org.junit.Ignore;
 import org.junit.runner.Description;
-import org.junit.runner.Runner;
 import org.junit.runner.notification.Failure;
 import org.junit.runner.notification.RunNotifier;
+import org.junit.runners.ParentRunner;
+import org.junit.runners.model.FrameworkMethod;
+import org.junit.runners.model.InitializationError;
 
 import java.lang.annotation.Annotation;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Type;
-import java.util.Arrays;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Random;
+import java.util.*;
+
+public class PropertyTestRunner extends ParentRunner<FrameworkMethod> {
 
-public class PropertyTestRunner extends Runner {
     private final Class<?> classUnderTest;
 
-    public PropertyTestRunner(Class<?> classUnderTest) {
+    public PropertyTestRunner(Class<?> classUnderTest) throws InitializationError{
+        super(classUnderTest);
         this.classUnderTest = classUnderTest;
     }
 
     @Override
-    public Description getDescription() {
-        return Description.createSuiteDescription(classUnderTest);
+    protected boolean isIgnored(FrameworkMethod child) {
+        return child.getAnnotation(Ignore.class) != null;
     }
 
-    private long getSeed(Method method) {
-        Seed seed = method.getAnnotation(Seed.class);
-        if (seed == null) {
-            return System.currentTimeMillis();
-        } else {
-            return seed.value();
+    @Override
+    protected List<FrameworkMethod> getChildren() {
+        List<FrameworkMethod> result = new ArrayList<>();
+        for (Method method : classUnderTest.getDeclaredMethods()) {
+            if (method.isAnnotationPresent(Property.class)
+                    && !method.isAnnotationPresent(Ignore.class)) {
+                result.add(new FrameworkMethod(method));
+            }
         }
+        return result;
     }
 
     @Override
-    public void run(RunNotifier notifier) {
-        Method[] methods = classUnderTest.getMethods();
-        for (Method method : methods) {
-            Property details = method.getAnnotation(Property.class);
-            if (details != null) {
-                Description description = Description.createTestDescription(classUnderTest, method.getName());
-                boolean failed = false;
-                int assumptionsFailed = 0;
-
-                long seed = getSeed(method);
-                Random random = new Random(seed);
-                int numRuns = details.runs();
-                for (int i = 0; i < numRuns && !failed; ++i) {
-                    int size = details.size();
-                    notifier.fireTestStarted(description);
-                    Object obj;
-                    try {
-                        obj = classUnderTest.getConstructor().newInstance();
-                    } catch (Throwable ex) { // HACKY
-                        System.out.println(ex);
-                        return;
-                    }
-                    RoseTree<Object[]> generated = generateArgs(random, size,
-                            method.getGenericParameterTypes(),
-                            method.getParameterAnnotations());
-                    try {
-                        method.invoke(obj, generated.getValue());
-                    } catch (InvocationTargetException ex) {
-                        if (ex.getTargetException() instanceof AssumptionViolatedException) {
-                            assumptionsFailed++;
+    protected Description describeChild(FrameworkMethod child) {
+        return Description.createTestDescription(classUnderTest, child.getName());
+    }
+
+    @Override
+    protected void runChild(FrameworkMethod child, RunNotifier notifier) {
+        try {
+            Property details = child.getAnnotation(Property.class);
+            Description description = Description.createTestDescription(classUnderTest, child.getName());
+            boolean failed = false;
+            int assumptionsFailed = 0;
+
+            long seed = getSeed(child.getMethod());
+            Random random = new Random(seed);
+            int numRuns = details.runs();
+            for (int i = 0; i < numRuns && !failed; ++i) {
+                int size = details.size();
+                notifier.fireTestStarted(description);
+                Object obj;
+                try {
+                    obj = classUnderTest.getConstructor().newInstance();
+                } catch (Throwable ex) { // HACKY
+                    System.out.println(ex);
+                    return;
+                }
+                RoseTree<Object[]> generated = generateArgs(random, size,
+                        child.getMethod().getGenericParameterTypes(),
+                        child.getMethod().getParameterAnnotations());
+                try {
+                    child.getMethod().invoke(obj, generated.getValue());
+                } catch (InvocationTargetException ex) {
+                    if (ex.getTargetException() instanceof AssumptionViolatedException) {
+                        assumptionsFailed++;
+                        i--;
+                    } else {
+                        System.out.println("Test failed with seed: " + seed);
+                        System.out.println("Failing arguments: " + Arrays.asList(generated.getValue()));
+                        Object[] shrinkResult = shrink(child.getMethod(), obj, generated);
+                        if (shrinkResult == null) {
+                            System.out.println("Arguments could not be shrunk any further");
                         } else {
-                            notifier.fireTestFailure(new Failure(description, ex.getTargetException()));
-                            System.out.println("Test failed with seed: " + seed);
-                            System.out.println("Failing arguments: " + Arrays.asList(generated.getValue()));
-                            Object[] shrinkResult = shrink(method, obj, generated);
-                            if (shrinkResult == null) {
-                                System.out.println("Arguments could not be shrunk any further");
-                            } else {
-                                System.out.println("Arguments shrunk to: " + Arrays.asList(shrinkResult));
-                            }
-                            failed = true;
+                            System.out.println("Arguments shrunk to: " + Arrays.asList(shrinkResult));
                         }
-                    } catch (IllegalAccessException ex) {
-                        notifier.fireTestFailure(new Failure(description, ex));
+                        notifier.fireTestFailure(new Failure(description, ex.getTargetException()));
                         failed = true;
                     }
+                } catch (IllegalAccessException ex) {
+                    notifier.fireTestFailure(new Failure(description, ex));
+                    failed = true;
                 }
+            }
 
-                if (assumptionsFailed > 0) {
-                    System.out.println("Failed " + assumptionsFailed + " assumptions");
-                }
-                notifier.fireTestFinished(description);
+            if (assumptionsFailed > 0) {
+                System.out.println("Failed " + assumptionsFailed + " assumptions");
             }
+            notifier.fireTestFinished(description);
+        } catch (Throwable ex) {
+            ex.printStackTrace();
+        }
+    }
+
+    private long getSeed(Method method) {
+        Seed seed = method.getAnnotation(Seed.class);
+        if (seed == null) {
+            return System.currentTimeMillis();
+        } else {
+            return seed.value();
         }
     }
 
@@ -103,12 +121,18 @@ public class PropertyTestRunner extends Runner {
             RoseTree<Object[]> tree = trees.next();
             try {
                 method.invoke(obj, tree.getValue());
-            } catch (Throwable ex) {
-                Iterator<RoseTree<Object[]>> children = tree.getChildren();
-                if (children.hasNext()) {
-                    trees = children;
+            } catch (InvocationTargetException ex) {
+                if (!(ex.getTargetException() instanceof AssumptionViolatedException)) {
+                    smallest = tree.getValue();
+                    Iterator<RoseTree<Object[]>> children = tree.getChildren();
+                    if (children.hasNext()) {
+                        trees = children;
+                    } else {
+                        break;
+                    }
                 }
-                smallest = tree.getValue();
+            } catch (IllegalAccessException ex) {
+                System.out.println(ex);
             }
         }
         return smallest;
@@ -135,14 +159,16 @@ public class PropertyTestRunner extends Runner {
                 generators[i] = getGeneratorFromType(types[i]);
 //            }
         }
-        return Generators.arrayGenerator(generators).generate(random, size);
+        @SuppressWarnings("unchecked")
+        Generator<Object>[] argsGenerators = (Generator<Object>[]) generators;
+        return Generator.tuple(argsGenerators).generate(random, size);
     }
 
     private Generator<?> getGeneratorFromType(Type type) {
         if (type instanceof Class) {
             Class<?> clazz = (Class<?>) type;
             if (clazz.isPrimitive() && clazz == Integer.TYPE) {
-                return Generators.integerGenerator();
+                return Generators.integer();
             } else {
                 throw new RuntimeException("Unknown type for generator (atm only int is supported)");
             }
diff --git a/src/main/java/au/id/zancanaro/RoseTree.java b/src/main/java/au/id/zancanaro/RoseTree.java
index 2a83ad2..bbcb949 100644
--- a/src/main/java/au/id/zancanaro/RoseTree.java
+++ b/src/main/java/au/id/zancanaro/RoseTree.java
@@ -4,6 +4,7 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 public class RoseTree<T> {
     private final T value;
@@ -36,14 +37,14 @@ public class RoseTree<T> {
 
     public static <T> Iterator<RoseTree<T>[]> permutations(RoseTree<T>[] trees) {
         return Iterators.flatten(Iterators.rangeIterator(trees.length, index ->
-                Iterators.mappingIterator(child -> {
-                    @SuppressWarnings("unchecked")
-                    RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length];
-                    for (int i = 0; i < trees.length; ++i) {
-                        result[i] = (i == index ? child : trees[i]);
-                    }
-                    return result;
-                }, trees[index].getChildren())
+                        Iterators.mappingIterator(child -> {
+                            @SuppressWarnings("unchecked")
+                            RoseTree<T>[] result = (RoseTree<T>[]) new RoseTree[trees.length];
+                            for (int i = 0; i < trees.length; ++i) {
+                                result[i] = (i == index ? child : trees[i]);
+                            }
+                            return result;
+                        }, trees[index].getChildren())
         ));
     }
 
@@ -69,4 +70,17 @@ public class RoseTree<T> {
     public <R> RoseTree<R> flatmap(Function<T, RoseTree<R>> f) {
         return RoseTree.join(this.fmap(f));
     }
+
+    public RoseTree<T> filter(Predicate<T> predicate) {
+        if (predicate.test(this.getValue())) {
+            return new RoseTree<>(
+                    this.getValue(),
+                    () -> Iterators.mappingIterator(tree -> tree.filter(predicate),
+                            Iterators.filteringIterator(
+                                    tree -> predicate.test(tree.getValue()),
+                                    this.getChildren())));
+        } else {
+            throw new IllegalArgumentException("Current value doesn't match predicate: whoops!");
+        }
+    }
 }
diff --git a/src/test/java/au/id/zancanaro/PropertyTests.java b/src/test/java/au/id/zancanaro/PropertyTests.java
index 0c9bb7e..3620e67 100644
--- a/src/test/java/au/id/zancanaro/PropertyTests.java
+++ b/src/test/java/au/id/zancanaro/PropertyTests.java
@@ -1,15 +1,28 @@
 package au.id.zancanaro;
 
 import au.id.zancanaro.annotations.Property;
-import au.id.zancanaro.annotations.Seed;
 import org.junit.Assert;
 import org.junit.Assume;
+import org.junit.experimental.theories.DataPoint;
 import org.junit.runner.RunWith;
 
 @RunWith(PropertyTestRunner.class)
 public class PropertyTests {
     @Property
-    public void method(int a, int b) {
+    public void aIsNotOdd(int a, int b) {
+        Assume.assumeFalse(a % 2 == 1);
+        Assert.assertFalse(a % 2 == 1);
+    }
+
+    @Property
+    public void aIsNotLessThanB(int a, int b) {
+        Assume.assumeFalse(a < b);
         Assert.assertFalse(a < b);
     }
+
+    @Property
+    public void aIsNotGreaterThanB(int a, int b) {
+//        Assume.assumeFalse(a > b);
+        Assert.assertFalse(a > b);
+    }
 }