2 files changed, 10 insertions, 27 deletions

diff --git a/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java b/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java
index 94032bc..a7bf19b 100644
--- a/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java
+++ b/src/main/java/au/id/zancanaro/javacheck/junit/Properties.java
@@ -1,11 +1,11 @@
 package au.id.zancanaro.javacheck.junit;
 
-import au.id.zancanaro.javacheck.DataSourceHelper;
 import au.id.zancanaro.javacheck.Generator;
 import au.id.zancanaro.javacheck.ShrinkResult;
 import au.id.zancanaro.javacheck.ShrinkTree;
 import au.id.zancanaro.javacheck.annotations.Property;
 import au.id.zancanaro.javacheck.annotations.Seed;
+import au.id.zancanaro.javacheck.object.GeneratorProvider;
 import org.junit.AssumptionViolatedException;
 import org.junit.runners.BlockJUnit4ClassRunner;
 import org.junit.runners.model.FrameworkMethod;
@@ -13,38 +13,24 @@ import org.junit.runners.model.InitializationError;
 import org.junit.runners.model.Statement;
 import org.junit.runners.model.TestClass;
 
+import java.lang.annotation.Annotation;
 import java.lang.reflect.Method;
-import java.lang.reflect.Type;
 import java.util.*;
 
 @SuppressWarnings("WeakerAccess")
 public class Properties extends BlockJUnit4ClassRunner {
-    private final DataSourceHelper helper;
+    private final GeneratorProvider provider;
 
     public Properties(Class<?> classObject) throws InitializationError {
         super(classObject);
-        helper = new DataSourceHelper(classObject);
+        provider = GeneratorProvider.DEFAULT_PROVIDER.withDataSources(classObject);
     }
 
     @Override
     protected void collectInitializationErrors(List<Throwable> errors) {
         super.collectInitializationErrors(errors);
-        Set<Type> generated = DataSourceHelper.validateGeneratorFields(getTestClass().getJavaClass(), errors);
-        validateTestMethodParameters(errors, generated);
     }
 
-    private void validateTestMethodParameters(List<Throwable> errors, Set<Type> generated) {
-        for (FrameworkMethod each : computeTestMethods()) {
-            for (Type type : each.getMethod().getGenericParameterTypes()) {
-                if (!generated.contains(type)) {
-                    errors.add(new Error("No @DataSource for type: " + type));
-                    generated.add(type); // ignore future errors on this type
-                }
-            }
-        }
-    }
-
-
     @Override
     protected void validateConstructor(List<Throwable> errors) {
         validateOnlyOneConstructor(errors);
@@ -73,18 +59,18 @@ public class Properties extends BlockJUnit4ClassRunner {
 
     @Override
     public Statement methodBlock(final FrameworkMethod method) {
-        return new GenerativeTester(method, getTestClass(), helper.computeGenerators());
+        return new GenerativeTester(method, getTestClass(), provider);
     }
 
     public static class GenerativeTester extends Statement {
         private final FrameworkMethod testMethod;
         private final TestClass testClass;
-        private final Map<Type, Generator<?>> generators;
+        private final GeneratorProvider provider;
 
-        public GenerativeTester(FrameworkMethod testMethod, TestClass testClass, Map<Type, Generator<?>> generators) {
+        public GenerativeTester(FrameworkMethod testMethod, TestClass testClass, GeneratorProvider provider) {
             this.testMethod = testMethod;
             this.testClass = testClass;
-            this.generators = generators;
+            this.provider = provider;
         }
 
         private static long getSeed(Method method) {
@@ -103,7 +89,7 @@ public class Properties extends BlockJUnit4ClassRunner {
                 runTest(new Object[0]);
             } else {
                 Generator<?>[] generators = Arrays.stream(method.getGenericParameterTypes())
-                        .map(this.generators::get)
+                        .map(param -> provider.getGenerator(param, new Annotation[0], provider))
                         .toArray(Generator[]::new);
                 Generator<List<Object>> generator = Generator.tuple(generators);
 
@@ -126,7 +112,6 @@ public class Properties extends BlockJUnit4ClassRunner {
                             throw new Error("Violated 50 assumptions in a row: failing test");
                         }
                     } catch (Throwable ex) {
-//                        tree.print(new OutputStreamWriter(System.out), Arrays::toString);
                         throw new PropertyError(method.getName(), seed, shrink(tree, ex));
                     }
                 }
@@ -178,8 +163,7 @@ public class Properties extends BlockJUnit4ClassRunner {
         public void runTest(final Object[] args) throws Throwable {
             new BlockJUnit4ClassRunner(testClass.getJavaClass()) {
                 @Override
-                protected void collectInitializationErrors(
-                        List<Throwable> errors) {
+                protected void collectInitializationErrors(List<Throwable> errors) {
                     // do nothing
                 }
 
diff --git a/src/main/java/au/id/zancanaro/javacheck/junit/PropertyError.java b/src/main/java/au/id/zancanaro/javacheck/junit/PropertyError.java
index d3747b9..5f89ca8 100644
--- a/src/main/java/au/id/zancanaro/javacheck/junit/PropertyError.java
+++ b/src/main/java/au/id/zancanaro/javacheck/junit/PropertyError.java
@@ -2,7 +2,6 @@ package au.id.zancanaro.javacheck.junit;
 
 import au.id.zancanaro.javacheck.ShrinkResult;
 
-import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;