1 files changed, 16 insertions, 8 deletions
diff --git a/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java b/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java
index 590b8d0..c513157 100644
--- a/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java
+++ b/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java
@@ -4,6 +4,8 @@ import io.trygvis.container.compiler.SqlUnitModel;
 import io.trygvis.container.compiler.model.Imports;
 import io.trygvis.container.compiler.model.MethodRef;
 import io.trygvis.container.compiler.model.TypeRef;
+import io.trygvis.persistence.generators.DaoGenerator;
+import io.trygvis.persistence.generators.DaoUtilsGenerator;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
@@ -20,6 +22,8 @@ import static org.testng.Assert.assertEquals;
 
 public class EntityMirrorTest {
 
+    GeneratorConfiguration generatorConfiguration = new GeneratorConfiguration();
+
     static TypeRef entityType = new TypeRef("Wat");
     static FieldMirror idLong = new FieldMirror(PRIMITIVE, FIELD, new TypeRef(Long.class), "id", "id", true, false, true);
     static FieldMirror idString = new FieldMirror(PRIMITIVE, FIELD, new TypeRef(String.class), "id", "id", true, false, true);
@@ -73,10 +77,11 @@ public class EntityMirrorTest {
         }
 
         SqlUnitModel unit = new SqlUnitModel().add(myTable).add(new SequenceMirror("seq-gen", "id_seq", 0, 0));
+        DaoGenerator generator = new DaoGenerator(generatorConfiguration, unit, myTable);
 
-        assertEquals(myTable.insertIntoSql(unit), insert);
-        assertEquals(myTable.deleteFromSql(), delete);
-        assertEquals(myTable.createTableSql(unit), create);
+        assertEquals(generator.insertIntoSql(), insert);
+        assertEquals(generator.deleteFromSql(), delete);
+        assertEquals(generator.createTableSql(), create);
     }
 
     @DataProvider(name = "insertIntoMethod", parallel = true)
@@ -115,10 +120,10 @@ public class EntityMirrorTest {
     }
 
     private MethodRef insertInto(FieldMirror... fields) {
-        EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
+        EntityMirror myTable = new EntityMirror(generatorConfiguration, entityType, "my_table");
         myTable.add(fields);
         SqlUnitModel unit = new SqlUnitModel().add(myTable);
-        return myTable.insertInto(unit, new Imports());
+        return new DaoUtilsGenerator(unit, generatorConfiguration, myTable).insertInto(new Imports());
     }
 
     @DataProvider(name = "deleteMethod", parallel = true)
@@ -141,7 +146,8 @@ public class EntityMirrorTest {
     private MethodRef delete(FieldMirror... fields) {
         EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
         myTable.add(fields);
-        return myTable.delete(new Imports());
+        SqlUnitModel unit = new SqlUnitModel().add(myTable);
+        return new DaoUtilsGenerator(unit, generatorConfiguration, myTable).delete(new Imports());
     }
 
     @DataProvider(name = "deleteByIdMethod", parallel = true)
@@ -170,7 +176,8 @@ public class EntityMirrorTest {
     private MethodRef deleteById(FieldMirror... fields) {
         EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
         myTable.add(fields);
-        return myTable.deleteById(new Imports());
+        SqlUnitModel unit = new SqlUnitModel().add(myTable);
+        return new DaoUtilsGenerator(unit, generatorConfiguration, myTable).deleteById(new Imports());
     }
 
     @Test
@@ -197,7 +204,8 @@ public class EntityMirrorTest {
     private MethodRef fromResultSet(FieldMirror... fields) {
         EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
         myTable.add(fields);
-        return myTable.fromResultSet(new Imports());
+        SqlUnitModel unit = new SqlUnitModel().add(myTable);
+        return new DaoUtilsGenerator(unit, generatorConfiguration, myTable).fromResultSet(new Imports());
     }
 
     private void eq(MethodRef m, String... expected) {