1 files changed, 10 insertions, 6 deletions
diff --git a/Swiften/FileTransfer/UnitTest/SOCKS5BytestreamServerSessionTest.cpp b/Swiften/FileTransfer/UnitTest/SOCKS5BytestreamServerSessionTest.cpp
index 06bc98f..cd480f0 100644
--- a/Swiften/FileTransfer/UnitTest/SOCKS5BytestreamServerSessionTest.cpp
+++ b/Swiften/FileTransfer/UnitTest/SOCKS5BytestreamServerSessionTest.cpp
@@ -28,25 +28,27 @@ class SOCKS5BytestreamServerSessionTest : public CppUnit::TestFixture {
 		CPPUNIT_TEST(testRequest_UnknownBytestream);
 		CPPUNIT_TEST(testReceiveData);
 		CPPUNIT_TEST(testReceiveData_Chunked);
 		CPPUNIT_TEST_SUITE_END();
 
 	public:
 		void setUp() {
 			receivedDataChunks = 0;
 			eventLoop = new DummyEventLoop();
+			bytestreams = new SOCKS5BytestreamRegistry();
 			connection = boost::shared_ptr<DummyConnection>(new DummyConnection(eventLoop));
 			connection->onDataSent.connect(boost::bind(&SOCKS5BytestreamServerSessionTest::handleDataWritten, this, _1));
 			stream1 = boost::shared_ptr<ByteArrayReadBytestream>(new ByteArrayReadBytestream(createByteArray("abcdefg")));
 		}
 
 		void tearDown() {
 			connection.reset();
+			delete bytestreams;
 			delete eventLoop;
 		}
 
 		void testAuthenticate() {
 			boost::shared_ptr<SOCKS5BytestreamServerSession> testling(createSession());
 			StartStopper<SOCKS5BytestreamServerSession> stopper(testling.get());
 
 			receive(createSafeByteArray("\x05\x02\x01\x02"));
 
@@ -61,19 +63,19 @@ class SOCKS5BytestreamServerSessionTest : public CppUnit::TestFixture {
 
 			CPPUNIT_ASSERT_EQUAL(0, static_cast<int>(receivedData.size()));
 			receive(createSafeByteArray("\x01"));
 			CPPUNIT_ASSERT(createByteArray("\x05\x00", 2) == receivedData);
 		}
 
 		void testRequest() {
 			boost::shared_ptr<SOCKS5BytestreamServerSession> testling(createSession());
 			StartStopper<SOCKS5BytestreamServerSession> stopper(testling.get());
-			bytestreams.addBytestream("abcdef", stream1);
+			bytestreams->addReadBytestream("abcdef", stream1);
 			authenticate();
 
 			ByteArray hostname(createByteArray("abcdef"));
 			receive(concat(createSafeByteArray("\x05\x01\x00\x03", 4), createSafeByteArray(hostname.size()), createSafeByteArray(hostname), createSafeByteArray("\x00\x00", 2)));
 			CPPUNIT_ASSERT(createByteArray("\x05\x00\x00\x03\x06\x61\x62\x63\x64\x65\x66\x00\x00", 13) == createByteArray(&receivedData[0], 13));
 		}
 
 		void testRequest_UnknownBytestream() {
 			boost::shared_ptr<SOCKS5BytestreamServerSession> testling(createSession());
@@ -82,37 +84,39 @@ class SOCKS5BytestreamServerSessionTest : public CppUnit::TestFixture {
 
 			ByteArray hostname(createByteArray("abcdef"));
 			receive(concat(createSafeByteArray("\x05\x01\x00\x03", 4), createSafeByteArray(hostname.size()), createSafeByteArray(hostname), createSafeByteArray("\x00\x00", 2)));
 			CPPUNIT_ASSERT(createByteArray("\x05\x04\x00\x03\x06\x61\x62\x63\x64\x65\x66\x00\x00", 13) == receivedData);
 		}
 
 		void testReceiveData() {
 			boost::shared_ptr<SOCKS5BytestreamServerSession> testling(createSession());
 			StartStopper<SOCKS5BytestreamServerSession> stopper(testling.get());
-			bytestreams.addBytestream("abcdef", stream1);
+			bytestreams->addReadBytestream("abcdef", stream1);
 			authenticate();
 			request("abcdef");
 			eventLoop->processEvents();
+			testling->startTransfer();
 			skipHeader("abcdef");
 
 			CPPUNIT_ASSERT(createByteArray("abcdefg") == receivedData);
 			CPPUNIT_ASSERT_EQUAL(2, receivedDataChunks);
 		}
 
 		void testReceiveData_Chunked() {
 			boost::shared_ptr<SOCKS5BytestreamServerSession> testling(createSession());
 			testling->setChunkSize(3);
 			StartStopper<SOCKS5BytestreamServerSession> stopper(testling.get());
-			bytestreams.addBytestream("abcdef", stream1);
+			bytestreams->addReadBytestream("abcdef", stream1);
 			authenticate();
 			request("abcdef");
 			eventLoop->processEvents();
-
+			testling->startTransfer();
+			eventLoop->processEvents();
 			skipHeader("abcdef");
 			CPPUNIT_ASSERT(createByteArray("abcdefg") == receivedData);
 			CPPUNIT_ASSERT_EQUAL(4, receivedDataChunks);
 		}
 
 	private:
 		void receive(const SafeByteArray& data) {
 			connection->receive(data);
 			eventLoop->processEvents();
@@ -135,23 +139,23 @@ class SOCKS5BytestreamServerSessionTest : public CppUnit::TestFixture {
 
 
 		void handleDataWritten(const SafeByteArray& data) {
 			receivedData.insert(receivedData.end(), data.begin(), data.end());
 			receivedDataChunks++;
 		}
 
 	private:
 		SOCKS5BytestreamServerSession* createSession() {
-			SOCKS5BytestreamServerSession* session = new SOCKS5BytestreamServerSession(connection, &bytestreams);
+			SOCKS5BytestreamServerSession* session = new SOCKS5BytestreamServerSession(connection, bytestreams);
 			return session;
 		}
 
 	private:
 		DummyEventLoop* eventLoop;
-		SOCKS5BytestreamRegistry bytestreams;
+		SOCKS5BytestreamRegistry* bytestreams;
 		boost::shared_ptr<DummyConnection> connection;
 		std::vector<unsigned char> receivedData;
 		int receivedDataChunks;
 		boost::shared_ptr<ByteArrayReadBytestream> stream1;
 };
 
 CPPUNIT_TEST_SUITE_REGISTRATION(SOCKS5BytestreamServerSessionTest);