Merge pull request #4 from mab/master

DataSetBuilder.addDataSet()
marcphilipp · Oct 29, 2013 · 15e879d · 15e879d
2 parents ba39cdb + 73b84d4
commit 15e879d
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 6 deletions.
diff --git a/src/main/java/org/dbunit/dataset/builder/DataSetBuilder.java b/src/main/java/org/dbunit/dataset/builder/DataSetBuilder.java
@@ -3,17 +3,14 @@
 import java.util.HashMap;
 import java.util.Map;
 
-import org.dbunit.dataset.CachedDataSet;
-import org.dbunit.dataset.DataSetException;
-import org.dbunit.dataset.IDataSet;
-import org.dbunit.dataset.ITableMetaData;
+import org.dbunit.dataset.*;
 import org.dbunit.dataset.stream.BufferedConsumer;
 import org.dbunit.dataset.stream.IDataSetConsumer;
 
 public class DataSetBuilder {
 
-	private final CachedDataSet dataSet = new CachedDataSet();
-	private final IDataSetConsumer consumer = new BufferedConsumer(dataSet);
+	private CachedDataSet dataSet = new CachedDataSet();
+	private IDataSetConsumer consumer = new BufferedConsumer(dataSet);
 	private final Map<String, TableMetaDataBuilder> tableNameToMetaData = new HashMap<String, TableMetaDataBuilder>();
 	private final StringPolicy stringPolicy;
 
@@ -53,6 +50,13 @@ public IDataSet build() throws DataSetException {
 		consumer.endDataSet();
 		return dataSet;
 	}
+
+	public void addDataSet(final IDataSet dataSet) throws DataSetException {
+		IDataSet[] dataSets = { build(), dataSet };
+		CompositeDataSet composite = new CompositeDataSet(dataSets);
+		this.dataSet = new CachedDataSet(composite);
+		consumer = new BufferedConsumer(this.dataSet);
+	}
 
 	protected void add(DataRowBuilder row) throws DataSetException {
 		ITableMetaData metaData = updateTableMetaData(row);

diff --git a/src/test/java/org/dbunit/dataset/builder/DataSetBuilderTest.java b/src/test/java/org/dbunit/dataset/builder/DataSetBuilderTest.java
@@ -106,4 +106,29 @@ public void addsNewColumnsOnTheFly() throws Exception {
 		assertNull(table.getValue(1, "NAME"));
 		assertEquals(18, table.getValue(1, "AGE"));
 	}
+
+	@Test
+	public void addDataSet() throws Exception {
+		DataSetBuilder builder = new DataSetBuilder();
+		builder.newRow("PERSON").with("NAME", "Bob").with("AGE", 18).add();
+
+		IDataSet dataSet = builder.build();
+		ITable table = dataSet.getTable("PERSON");
+
+		assertEquals(1, table.getRowCount());
+
+		builder = new DataSetBuilder();
+		builder.newRow("PERSON").with("NAME", "John").with("AGE", 19).add();
+		builder.addDataSet(dataSet);
+
+		IDataSet dataSet2 = builder.build();
+	   	ITable table2 = dataSet2.getTable("PERSON");
+
+	   	assertEquals(2, table2.getRowCount());
+	   	assertEquals("John", table2.getValue(0, "NAME"));
+	   	assertEquals(19, table2.getValue(0, "AGE"));
+
+	   	assertEquals("Bob", table2.getValue(1, "NAME"));
+	   	assertEquals(18, table2.getValue(1, "AGE"));
+	}
 }