Builders for incoming and outgoing headers and batches. The protocol and serial components can now be used to serialise, sign, deserialise and verify real bundles (except for message parsing).

test/build.xml

@@ -21,6 +21,7 @@
 			<test name='net.sf.briar.i18n.I18nTest'/>
 			<test name='net.sf.briar.invitation.InvitationWorkerTest'/>
 			<test name='net.sf.briar.protocol.BundleReaderTest'/>
+			<test name='net.sf.briar.protocol.BundleReadWriteTest'/>
 			<test name='net.sf.briar.protocol.BundleWriterTest'/>
 			<test name='net.sf.briar.serial.ReaderImplTest'/>
 			<test name='net.sf.briar.serial.WriterImplTest'/>

test/net/sf/briar/db/DatabaseComponentTest.java

@@ -1,7 +1,5 @@
 package net.sf.briar.db;
 
-import java.io.IOException;
-import java.security.SignatureException;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
@@ -18,9 +16,9 @@ import net.sf.briar.api.protocol.AuthorId;
 import net.sf.briar.api.protocol.Batch;
 import net.sf.briar.api.protocol.BatchBuilder;
 import net.sf.briar.api.protocol.BatchId;
-import net.sf.briar.api.protocol.Bundle;
-import net.sf.briar.api.protocol.BundleBuilder;
 import net.sf.briar.api.protocol.BundleId;
+import net.sf.briar.api.protocol.BundleReader;
+import net.sf.briar.api.protocol.BundleWriter;
 import net.sf.briar.api.protocol.GroupId;
 import net.sf.briar.api.protocol.Header;
 import net.sf.briar.api.protocol.HeaderBuilder;
@@ -507,7 +505,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 
 	@Test
 	public void testGenerateBundleThrowsExceptionIfContactIsMissing()
-	throws DbException, IOException, SignatureException {
+	throws Exception {
 		Mockery context = new Mockery();
 		@SuppressWarnings("unchecked")
 		final Database<Object> database = context.mock(Database.class);
@@ -518,7 +516,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 		@SuppressWarnings("unchecked")
 		final Provider<BatchBuilder> batchBuilderProvider =
 			context.mock(Provider.class, "batchBuilderProvider");
-		final BundleBuilder bundleBuilder = context.mock(BundleBuilder.class);
+		final BundleWriter bundleBuilder = context.mock(BundleWriter.class);
 		context.checking(new Expectations() {{
 			// Check that the contact is still in the DB
 			oneOf(database).startTransaction();
@@ -539,8 +537,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 	}
 
 	@Test
-	public void testGenerateBundle() throws DbException, IOException,
-	SignatureException {
+	public void testGenerateBundle() throws Exception {
 		final long headerSize = 1234L;
 		Mockery context = new Mockery();
 		@SuppressWarnings("unchecked")
@@ -552,12 +549,11 @@ public abstract class DatabaseComponentTest extends TestCase {
 		@SuppressWarnings("unchecked")
 		final Provider<BatchBuilder> batchBuilderProvider =
 			context.mock(Provider.class, "batchBuilderProvider");
-		final BundleBuilder bundleBuilder = context.mock(BundleBuilder.class);
+		final BundleWriter bundleWriter = context.mock(BundleWriter.class);
 		final HeaderBuilder headerBuilder = context.mock(HeaderBuilder.class);
 		final Header header = context.mock(Header.class);
 		final BatchBuilder batchBuilder = context.mock(BatchBuilder.class);
 		final Batch batch = context.mock(Batch.class);
-		final Bundle bundle = context.mock(Bundle.class);
 		context.checking(new Expectations() {{
 			allowing(database).startTransaction();
 			will(returnValue(txn));
@@ -582,11 +578,11 @@ public abstract class DatabaseComponentTest extends TestCase {
 			// Build the header
 			oneOf(headerBuilder).build();
 			will(returnValue(header));
-			oneOf(bundleBuilder).getCapacity();
+			oneOf(bundleWriter).getCapacity();
 			will(returnValue(1024L * 1024L));
 			oneOf(header).getSize();
 			will(returnValue(headerSize));
-			oneOf(bundleBuilder).addHeader(header);
+			oneOf(bundleWriter).addHeader(header);
 			// Add a batch to the bundle
 			oneOf(database).getSendableMessages(txn, contactId,
 					Batch.MAX_SIZE - headerSize);
@@ -604,25 +600,24 @@ public abstract class DatabaseComponentTest extends TestCase {
 			oneOf(database).addOutstandingBatch(
 					txn, contactId, batchId, messages);
 			// Add the batch to the bundle
-			oneOf(bundleBuilder).addBatch(batch);
+			oneOf(bundleWriter).addBatch(batch);
 			// Check whether to add another batch
 			oneOf(batch).getSize();
 			will(returnValue((long) message.getSize()));
 			// No, just send the bundle
-			oneOf(bundleBuilder).build();
-			will(returnValue(bundle));
+			oneOf(bundleWriter).close();
 		}});
 		DatabaseComponent db = createDatabaseComponent(database, cleaner,
 				headerBuilderProvider, batchBuilderProvider);
 
-		db.generateBundle(contactId, bundleBuilder);
+		db.generateBundle(contactId, bundleWriter);
 
 		context.assertIsSatisfied();
 	}
 
 	@Test
 	public void testReceiveBundleThrowsExceptionIfContactIsMissing()
-	throws DbException, IOException, SignatureException {
+	throws Exception {
 		Mockery context = new Mockery();
 		@SuppressWarnings("unchecked")
 		final Database<Object> database = context.mock(Database.class);
@@ -633,7 +628,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 		@SuppressWarnings("unchecked")
 		final Provider<BatchBuilder> batchBuilderProvider =
 			context.mock(Provider.class, "batchBuilderProvider");
-		final Bundle bundle = context.mock(Bundle.class);
+		final BundleReader bundleReader = context.mock(BundleReader.class);
 		context.checking(new Expectations() {{
 			// Check that the contact is still in the DB
 			oneOf(database).startTransaction();
@@ -646,7 +641,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 				headerBuilderProvider, batchBuilderProvider);
 
 		try {
-			db.receiveBundle(contactId, bundle);
+			db.receiveBundle(contactId, bundleReader);
 			assertTrue(false);
 		} catch(NoSuchContactException expected) {}
 
@@ -654,8 +649,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 	}
 
 	@Test
-	public void testReceivedBundle() throws DbException, IOException,
-	SignatureException {
+	public void testReceiveBundle() throws Exception {
 		Mockery context = new Mockery();
 		@SuppressWarnings("unchecked")
 		final Database<Object> database = context.mock(Database.class);
@@ -666,7 +660,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 		@SuppressWarnings("unchecked")
 		final Provider<BatchBuilder> batchBuilderProvider =
 			context.mock(Provider.class, "batchBuilderProvider");
-		final Bundle bundle = context.mock(Bundle.class);
+		final BundleReader bundleReader = context.mock(BundleReader.class);
 		final Header header = context.mock(Header.class);
 		final Batch batch = context.mock(Batch.class);
 		context.checking(new Expectations() {{
@@ -676,7 +670,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 			allowing(database).containsContact(txn, contactId);
 			will(returnValue(true));
 			// Header
-			oneOf(bundle).getHeader();
+			oneOf(bundleReader).getHeader();
 			will(returnValue(header));
 			// Acks
 			oneOf(header).getAcks();
@@ -692,7 +686,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 			will(returnValue(transports));
 			oneOf(database).setTransports(txn, contactId, transports);
 			// Batches
-			oneOf(bundle).getNextBatch();
+			oneOf(bundleReader).getNextBatch();
 			will(returnValue(batch));
 			oneOf(batch).getMessages();
 			will(returnValue(Collections.singleton(message)));
@@ -706,8 +700,9 @@ public abstract class DatabaseComponentTest extends TestCase {
 			will(returnValue(batchId));
 			oneOf(database).addBatchToAck(txn, contactId, batchId);
 			// Any more batches? Nope
-			oneOf(bundle).getNextBatch();
+			oneOf(bundleReader).getNextBatch();
 			will(returnValue(null));
+			oneOf(bundleReader).close();
 			// Lost batches
 			oneOf(header).getId();
 			will(returnValue(bundleId));
@@ -718,7 +713,7 @@ public abstract class DatabaseComponentTest extends TestCase {
 		DatabaseComponent db = createDatabaseComponent(database, cleaner,
 				headerBuilderProvider, batchBuilderProvider);
 
-		db.receiveBundle(contactId, bundle);
+		db.receiveBundle(contactId, bundleReader);
 
 		context.assertIsSatisfied();
 	}

test/net/sf/briar/db/H2DatabaseTest.java

@@ -110,7 +110,7 @@ public class H2DatabaseTest extends TestCase {
 		assertEquals(authorId, m1.getAuthor());
 		assertEquals(timestamp, m1.getTimestamp());
 		assertEquals(size, m1.getSize());
-		assertTrue(Arrays.equals(body, m1.getBody()));
+		assertTrue(Arrays.equals(body, m1.getBytes()));
 		// Delete the records
 		db.removeContact(txn, contactId);
 		db.removeMessage(txn, messageId);

test/net/sf/briar/protocol/BundleReadWriteTest.java

@@ -0,0 +1,156 @@
+package net.sf.briar.protocol;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.security.KeyPair;
+import java.security.KeyPairGenerator;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.security.Signature;
+import java.security.SignatureException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+import junit.framework.TestCase;
+import net.sf.briar.TestUtils;
+import net.sf.briar.api.protocol.AuthorId;
+import net.sf.briar.api.protocol.Batch;
+import net.sf.briar.api.protocol.BatchBuilder;
+import net.sf.briar.api.protocol.BatchId;
+import net.sf.briar.api.protocol.BundleReader;
+import net.sf.briar.api.protocol.BundleWriter;
+import net.sf.briar.api.protocol.GroupId;
+import net.sf.briar.api.protocol.Header;
+import net.sf.briar.api.protocol.HeaderBuilder;
+import net.sf.briar.api.protocol.Message;
+import net.sf.briar.api.protocol.MessageId;
+import net.sf.briar.api.protocol.MessageParser;
+import net.sf.briar.api.protocol.UniqueId;
+import net.sf.briar.api.serial.FormatException;
+import net.sf.briar.api.serial.Reader;
+import net.sf.briar.api.serial.Writer;
+import net.sf.briar.api.serial.WriterFactory;
+import net.sf.briar.serial.ReaderFactoryImpl;
+import net.sf.briar.serial.WriterFactoryImpl;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import com.google.inject.Provider;
+
+public class BundleReadWriteTest extends TestCase {
+
+	private static final String SIGNATURE_ALGO = "SHA256withRSA";
+	private static final String KEY_PAIR_ALGO = "RSA";
+	private static final String DIGEST_ALGO = "SHA-256";
+
+	private final File testDir = TestUtils.getTestDirectory();
+	private final File bundle = new File(testDir, "bundle");
+
+	private final long capacity = 1024L;
+	private final BatchId ack = new BatchId(TestUtils.getRandomId());
+	private final Set<BatchId> acks = Collections.singleton(ack);
+	private final GroupId sub = new GroupId(TestUtils.getRandomId());
+	private final Set<GroupId> subs = Collections.singleton(sub);
+	private final Map<String, String> transports =
+		Collections.singletonMap("foo", "bar");
+
+	private final MessageId messageId = new MessageId(TestUtils.getRandomId());
+	private final AuthorId authorId = new AuthorId(TestUtils.getRandomId());
+	private final long timestamp = System.currentTimeMillis();
+	private final byte[] messageBody = new byte[123];
+	private final Message message = new MessageImpl(messageId, MessageId.NONE,
+			sub, authorId, timestamp, messageBody);
+
+	// FIXME: This test should not depend on an impl in another component
+	private final WriterFactory wf = new WriterFactoryImpl();
+
+	private final KeyPair keyPair;
+	private final Signature sig;
+	private final MessageDigest digest;
+
+	public BundleReadWriteTest() throws NoSuchAlgorithmException {
+		super();
+		keyPair = KeyPairGenerator.getInstance(KEY_PAIR_ALGO).generateKeyPair();
+		sig = Signature.getInstance(SIGNATURE_ALGO);
+		digest = MessageDigest.getInstance(DIGEST_ALGO);
+		assertEquals(digest.getDigestLength(), UniqueId.LENGTH);
+	}
+
+	@Before
+	public void setUp() {
+		testDir.mkdirs();
+	}
+
+	@Test
+	public void testWriteBundle() throws Exception {
+		HeaderBuilder h = new OutgoingHeaderBuilder(keyPair, sig, digest, wf);
+		h.addAcks(acks);
+		h.addSubscriptions(subs);
+		h.addTransports(transports);
+		Header header = h.build();
+
+		BatchBuilder b = new OutgoingBatchBuilder(keyPair, sig, digest, wf);
+		b.addMessage(message);
+		Batch batch = b.build();
+
+		FileOutputStream out = new FileOutputStream(bundle);
+		Writer writer = new WriterFactoryImpl().createWriter(out);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
+
+		w.addHeader(header);
+		w.addBatch(batch);
+		w.close();
+
+		assertTrue(bundle.exists());
+		assertTrue(bundle.length() > messageBody.length);
+	}
+
+	@Test
+	public void testWriteAndReadBundle() throws Exception {
+
+		testWriteBundle();
+
+		MessageParser messageParser = new MessageParser() {
+			public Message parseMessage(byte[] body) throws FormatException,
+			SignatureException {
+				// FIXME: Really parse the message
+				return message;
+			}
+		};
+		Provider<HeaderBuilder> headerBuilderProvider =
+			new Provider<HeaderBuilder>() {
+			public HeaderBuilder get() {
+				return new IncomingHeaderBuilder(keyPair, sig, digest, wf);
+			}
+		};
+		Provider<BatchBuilder> batchBuilderProvider =
+			new Provider<BatchBuilder>() {
+			public BatchBuilder get() {
+				return new IncomingBatchBuilder(keyPair, sig, digest, wf);
+			}
+		};
+
+		FileInputStream in = new FileInputStream(bundle);
+		Reader reader = new ReaderFactoryImpl().createReader(in);
+		BundleReader r = new BundleReaderImpl(reader, bundle.length(),
+				messageParser, headerBuilderProvider, batchBuilderProvider);
+
+		Header h = r.getHeader();
+		assertEquals(acks, h.getAcks());
+		assertEquals(subs, h.getSubscriptions());
+		assertEquals(transports, h.getTransports());
+		Batch b = r.getNextBatch();
+		assertEquals(Collections.singletonList(message), b.getMessages());
+		assertNull(r.getNextBatch());
+		r.close();
+	}
+
+	@After
+	public void tearDown() {
+		TestUtils.deleteTestDirectory(testDir);
+	}
+}

test/net/sf/briar/protocol/BundleReaderTest.java

@@ -12,6 +12,7 @@ import net.sf.briar.TestUtils;
 import net.sf.briar.api.protocol.Batch;
 import net.sf.briar.api.protocol.BatchBuilder;
 import net.sf.briar.api.protocol.BatchId;
+import net.sf.briar.api.protocol.BundleReader;
 import net.sf.briar.api.protocol.GroupId;
 import net.sf.briar.api.protocol.Header;
 import net.sf.briar.api.protocol.HeaderBuilder;
@@ -46,7 +47,7 @@ public class BundleReaderTest extends TestCase {
 	private final byte[] batchSig = TestUtils.getRandomId();
 
 	@Test
-	public void testGetHeader() throws IOException, SignatureException {
+	public void testGetHeader() throws Exception {
 		Mockery context = new Mockery();
 		final Reader reader = context.mock(Reader.class);
 		final MessageParser messageParser = context.mock(MessageParser.class);
@@ -82,7 +83,7 @@ public class BundleReaderTest extends TestCase {
 			oneOf(headerBuilder).build();
 			will(returnValue(header));
 		}});
-		BundleReader r = createBundleReader(reader, messageParser,
+		BundleReader r = new BundleReaderImpl(reader, size, messageParser,
 				headerBuilderProvider, batchBuilderProvider);
 
 		assertEquals(header, r.getHeader());
@@ -91,8 +92,7 @@ public class BundleReaderTest extends TestCase {
 	}
 
 	@Test
-	public void testBatchBeforeHeaderThrowsException() throws IOException,
-	SignatureException {
+	public void testBatchBeforeHeaderThrowsException() throws Exception {
 		Mockery context = new Mockery();
 		final Reader reader = context.mock(Reader.class);
 		final MessageParser messageParser = context.mock(MessageParser.class);
@@ -102,7 +102,7 @@ public class BundleReaderTest extends TestCase {
 		@SuppressWarnings("unchecked")
 		final Provider<BatchBuilder> batchBuilderProvider =
 			context.mock(Provider.class, "batchBuilderProvider");
-		BundleReader r = createBundleReader(reader, messageParser,
+		BundleReader r = new BundleReaderImpl(reader, size, messageParser,
 				headerBuilderProvider, batchBuilderProvider);
 
 		try {
@@ -114,8 +114,30 @@ public class BundleReaderTest extends TestCase {
 	}
 
 	@Test
-	public void testGetHeaderNoBatches() throws IOException,
+	public void testCloseBeforeHeaderDoesNotThrowException() throws IOException,
 	SignatureException {
+		Mockery context = new Mockery();
+		final Reader reader = context.mock(Reader.class);
+		final MessageParser messageParser = context.mock(MessageParser.class);
+		@SuppressWarnings("unchecked")
+		final Provider<HeaderBuilder> headerBuilderProvider =
+			context.mock(Provider.class);
+		@SuppressWarnings("unchecked")
+		final Provider<BatchBuilder> batchBuilderProvider =
+			context.mock(Provider.class, "batchBuilderProvider");
+		context.checking(new Expectations() {{
+			oneOf(reader).close();
+		}});
+		BundleReader r = new BundleReaderImpl(reader, size, messageParser,
+				headerBuilderProvider, batchBuilderProvider);
+
+		r.close();
+
+		context.assertIsSatisfied();
+	}
+
+	@Test
+	public void testGetHeaderNoBatches() throws Exception {
 		Mockery context = new Mockery();
 		final Reader reader = context.mock(Reader.class);
 		final MessageParser messageParser = context.mock(MessageParser.class);
@@ -155,19 +177,21 @@ public class BundleReaderTest extends TestCase {
 			oneOf(reader).hasListEnd();
 			will(returnValue(true));
 			oneOf(reader).readListEnd();
+			// Close
+			oneOf(reader).close();
 		}});
-		BundleReader r = createBundleReader(reader, messageParser,
+		BundleReader r = new BundleReaderImpl(reader, size, messageParser,
 				headerBuilderProvider, batchBuilderProvider);
 
 		assertEquals(header, r.getHeader());
 		assertNull(r.getNextBatch());
+		r.close();
 
 		context.assertIsSatisfied();
 	}
 
 	@Test
-	public void testGetHeaderOneBatch() throws IOException,
-	SignatureException {
+	public void testGetHeaderOneBatch() throws Exception {
 		Mockery context = new Mockery();
 		final Reader reader = context.mock(Reader.class);
 		final MessageParser messageParser = context.mock(MessageParser.class);
@@ -226,26 +250,17 @@ public class BundleReaderTest extends TestCase {
 			oneOf(reader).hasListEnd();
 			will(returnValue(true));
 			oneOf(reader).readListEnd();
+			// Close
+			oneOf(reader).close();
 		}});
-		BundleReader r = createBundleReader(reader, messageParser,
+		BundleReader r = new BundleReaderImpl(reader, size, messageParser,
 				headerBuilderProvider, batchBuilderProvider);
 
 		assertEquals(header, r.getHeader());
 		assertEquals(batch, r.getNextBatch());
 		assertNull(r.getNextBatch());
+		r.close();
 
 		context.assertIsSatisfied();
 	}
-
-	private BundleReader createBundleReader(Reader reader,
-			MessageParser messageParser,
-			Provider<HeaderBuilder> headerBuilderProvider,
-			Provider<BatchBuilder> batchBuilderProvider) {
-		return new BundleReader(reader, messageParser, headerBuilderProvider,
-				batchBuilderProvider) {
-			public long getSize() {
-				return size;
-			}
-		};
-	}
 }

test/net/sf/briar/protocol/BundleWriterTest.java

@@ -9,7 +9,7 @@ import junit.framework.TestCase;
 import net.sf.briar.TestUtils;
 import net.sf.briar.api.protocol.Batch;
 import net.sf.briar.api.protocol.BatchId;
-import net.sf.briar.api.protocol.Bundle;
+import net.sf.briar.api.protocol.BundleWriter;
 import net.sf.briar.api.protocol.GroupId;
 import net.sf.briar.api.protocol.Header;
 import net.sf.briar.api.protocol.Message;
@@ -21,6 +21,7 @@ import org.junit.Test;
 
 public class BundleWriterTest extends TestCase {
 
+	private final long capacity = 1024L * 1024L;
 	private final BatchId ack = new BatchId(TestUtils.getRandomId());
 	private final Set<BatchId> acks = Collections.singleton(ack);
 	private final GroupId sub = new GroupId(TestUtils.getRandomId());
@@ -28,7 +29,6 @@ public class BundleWriterTest extends TestCase {
 	private final Map<String, String> transports =
 		Collections.singletonMap("foo", "bar");
 	private final byte[] headerSig = TestUtils.getRandomId();
-	private final long capacity = 1024L * 1024L;
 	private final byte[] messageBody = new byte[123];
 	private final byte[] batchSig = TestUtils.getRandomId();
 
@@ -59,7 +59,7 @@ public class BundleWriterTest extends TestCase {
 			will(returnValue(headerSig));
 			oneOf(writer).writeRaw(headerSig);
 		}});
-		BundleWriter w = createBundleWriter(writer);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
 
 		w.addHeader(header);
 
@@ -91,7 +91,7 @@ public class BundleWriterTest extends TestCase {
 			will(returnValue(headerSig));
 			oneOf(writer).writeRaw(headerSig);
 		}});
-		BundleWriter w = createBundleWriter(writer);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
 
 		w.addHeader(header);
 
@@ -103,7 +103,7 @@ public class BundleWriterTest extends TestCase {
 		Mockery context = new Mockery();
 		final Writer writer = context.mock(Writer.class);
 		final Batch batch = context.mock(Batch.class);
-		BundleWriter w = createBundleWriter(writer);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
 
 		try {
 			w.addBatch(batch);
@@ -117,7 +117,7 @@ public class BundleWriterTest extends TestCase {
 	public void testCloseBeforeHeaderThrowsException() throws IOException {
 		Mockery context = new Mockery();
 		final Writer writer = context.mock(Writer.class);
-		BundleWriter w = createBundleWriter(writer);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
 
 		try {
 			w.close();
@@ -159,7 +159,7 @@ public class BundleWriterTest extends TestCase {
 			oneOf(writer).writeListEnd();
 			oneOf(writer).close();
 		}});
-		BundleWriter w = createBundleWriter(writer);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
 
 		w.addHeader(header);
 		w.close();
@@ -200,7 +200,7 @@ public class BundleWriterTest extends TestCase {
 			oneOf(writer).writeListStart();
 			oneOf(batch).getMessages();
 			will(returnValue(Collections.singleton(message)));
-			oneOf(message).getBody();
+			oneOf(message).getBytes();
 			will(returnValue(messageBody));
 			oneOf(writer).writeRaw(messageBody);
 			oneOf(writer).writeListEnd();
@@ -211,7 +211,7 @@ public class BundleWriterTest extends TestCase {
 			oneOf(writer).writeListStart();
 			oneOf(batch).getMessages();
 			will(returnValue(Collections.singleton(message)));
-			oneOf(message).getBody();
+			oneOf(message).getBytes();
 			will(returnValue(messageBody));
 			oneOf(writer).writeRaw(messageBody);
 			oneOf(writer).writeListEnd();
@@ -222,7 +222,7 @@ public class BundleWriterTest extends TestCase {
 			oneOf(writer).writeListEnd();
 			oneOf(writer).close();
 		}});
-		BundleWriter w = createBundleWriter(writer);
+		BundleWriter w = new BundleWriterImpl(writer, capacity);
 
 		w.addHeader(header);
 		w.addBatch(batch);
@@ -231,12 +231,4 @@ public class BundleWriterTest extends TestCase {
 
 		context.assertIsSatisfied();
 	}
-
-	private BundleWriter createBundleWriter(Writer writer) {
-		return new BundleWriter(writer, capacity) {
-			public Bundle build() throws IOException {
-				return null;
-			}
-		};
-	}
 }