2 files changed, 26 insertions, 17 deletions

diff --git a/p2pvr/daemon/unittests/mockDatasource.cpp b/p2pvr/daemon/unittests/mockDatasource.cpp
index da4e3c6..cb3c95a 100644
--- a/p2pvr/daemon/unittests/mockDatasource.cpp
+++ b/p2pvr/daemon/unittests/mockDatasource.cpp
@@ -6,22 +6,21 @@
 #include <boost/filesystem/operations.hpp>
 #include <definedDirs.h>
 
-typedef MockDatabase::MockConnectionLoader MDMCL;
-DECLARE_CUSTOM_COMPONENT_LOADER("mock", T, MDMCL, ConnectionLoader);
+DECLARE_CUSTOM_COMPONENT_LOADER("mock", T, MockConnectionLoader, ConnectionLoader);
 
-std::map<std::string, std::string> MockDatabase::mocks;
-unsigned int MockDatabase::mocked = 0;
+std::map<std::string, MockConnectionLoader::ConnectionCreator> MockConnectionLoader::mocks;
+unsigned int MockConnectionLoader::mocked = 0;
 
 DB::Connection *
-MockDatabase::MockConnectionLoader::create(const std::string & n) const
+MockConnectionLoader::create(const std::string & n) const
 {
-	auto m = MockDatabase::mocks.find(n);
-	return InstanceMap<ConnectionLoader, std::string>::Get<std::invalid_argument>("postgresql")->create(m->second);
+	return mocks.find(n)->second();
 }
 
 MockDatabase::MockDatabase(const std::string & masterdb, const std::string & name, const std::vector<std::string> & ss) :
 	master(InstanceMap<ConnectionLoader, std::string>::Get<std::invalid_argument>("postgresql")->create(masterdb)),
-	testDbName(stringbf("test_%d_%d", getpid(), ++mocked))
+	testDbName(stringbf("test_%d_%d", getpid(), ++MockConnectionLoader::mocked)),
+	mockName(name)
 {
 	Logger()->messagebf(LOG_DEBUG, "Setting up new mocked database %s", testDbName);
 	DropDatabase();
@@ -32,7 +31,7 @@ MockDatabase::MockDatabase(const std::string & masterdb, const std::string & nam
 			system(("psql -q -1 -U postgres " + testDbName + " -f " + (RootDir  / s).string()).c_str());
 		}
 		Logger()->messagebf(LOG_DEBUG, "%s initialized", testDbName);
-		mocks[name] = stringbf("user=postgres dbname=%s", testDbName);
+		MockConnectionLoader::mocks[name] = boost::bind(MockDatabase::openConnection, stringbf("user=postgres dbname=%s", testDbName));
 	}
 	catch (...) {
 		DropDatabase();
@@ -40,11 +39,18 @@ MockDatabase::MockDatabase(const std::string & masterdb, const std::string & nam
 	}
 }
 
+DB::Connection *
+MockDatabase::openConnection(const std::string & connStr)
+{
+	return InstanceMap<ConnectionLoader, std::string>::Get<std::invalid_argument>("postgresql")->create(connStr);
+}
+
 MockDatabase::~MockDatabase()
 {
 	Logger()->messagebf(LOG_DEBUG, "Tearing down mocked database %s", testDbName);
 	DropDatabase();
 	delete master;
+	MockConnectionLoader::mocks.erase(mockName);
 	Logger()->messagebf(LOG_DEBUG, "%s torn down", testDbName);
 }
 
diff --git a/p2pvr/daemon/unittests/mockDatasource.h b/p2pvr/daemon/unittests/mockDatasource.h
index c39e560..df68602 100644
--- a/p2pvr/daemon/unittests/mockDatasource.h
+++ b/p2pvr/daemon/unittests/mockDatasource.h
@@ -3,13 +3,17 @@
 
 #include <rdbmsDataSource.h>
 
-class MockDatabase {
+class MockConnectionLoader : public ConnectionLoader {
 	public:
-		class MockConnectionLoader : public ConnectionLoader {
-			public:
-				DB::Connection * create(const std::string &) const;
-		};
+		DB::Connection * create(const std::string &) const;
+
+		static unsigned int mocked;
+		typedef boost::function<DB::Connection *()> ConnectionCreator;
+		static std::map<std::string, ConnectionCreator> mocks;
+};
 
+class MockDatabase {
+	public:
 		MockDatabase(const std::string & master, const std::string & name, const std::vector<std::string> & ss);
 		~MockDatabase();
 
@@ -18,11 +22,10 @@ class MockDatabase {
 		void CreateNewDatabase() const;
 
 	private:
+		static DB::Connection * openConnection(const std::string & connStr);
 		const DB::Connection * master;
 		const std::string testDbName;
-
-		static unsigned int mocked;
-		static std::map<std::string, std::string> mocks;
+		const std::string mockName;
 };
 
 #endif