3 files changed, 29 insertions, 7 deletions
diff --git a/Swift/Controllers/UnitTest/RosterControllerTest.cpp b/Swift/Controllers/UnitTest/RosterControllerTest.cpp
index 84a1063..0296cdc 100644
--- a/Swift/Controllers/UnitTest/RosterControllerTest.cpp
+++ b/Swift/Controllers/UnitTest/RosterControllerTest.cpp
@@ -53,7 +53,7 @@ class RosterControllerTest : public CppUnit::TestFixture
 			groups.push_back("testGroup2");
 			xmppRoster_->addContact(JID("test@testdomain.com/bob"), "name", groups);
 			
-			//CPPUNIT_ASSERT(treeWidgetFactory_->);
+			CPPUNIT_ASSERT_EQUAL(2, (int)treeWidgetFactory_->getGroups().size());
 			//CPPUNIT_ASSERT_EQUAL(String("Bob"), xmppRoster_->getNameForJID(JID("foo@bar.com")));
 		};
 
diff --git a/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h b/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h
index 09e4742..b2b4f10 100644
--- a/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h
+++ b/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h
@@ -2,6 +2,9 @@
 #define SWIFTEN_MockTreeWidgetFactory_H
 
 #include "Swiften/Roster/TreeWidgetFactory.h"
+
+#include <vector>
+#include "Swiften/Base/foreach.h"
 #include "Swiften/Roster/UnitTest/MockTreeWidget.h"
 #include "Swiften/Roster/UnitTest/MockTreeWidgetItem.h"
 
@@ -14,14 +17,30 @@ class MockTreeWidgetFactory : public TreeWidgetFactory {
 	public:
 		virtual ~MockTreeWidgetFactory() {}
 		virtual TreeWidget* createTreeWidget() {
-			return new MockTreeWidget();
+			root_ = new MockTreeWidget();
+			return root_;
 		};
-		virtual TreeWidgetItem* createTreeWidgetItem(TreeWidgetItem*) {
-			return new MockTreeWidgetItem();
+		virtual TreeWidgetItem* createTreeWidgetItem(TreeWidgetItem* group) {
+			MockTreeWidgetItem* entry = new MockTreeWidgetItem();
+			groupMembers_[group].push_back(entry);
+			return entry;
 		};
 		virtual TreeWidgetItem* createTreeWidgetItem(TreeWidget*) {
-			return new MockTreeWidgetItem();
-		}
+			MockTreeWidgetItem* group = new MockTreeWidgetItem();
+			groups_.push_back(group);
+			return group;
+		};
+		virtual std::vector<String> getGroups() {
+			std::vector<String> groupNames;
+			foreach (MockTreeWidgetItem* group, groups_) {
+				groupNames.push_back(group->getText());
+			}
+			return groupNames;
+		};
+	private:
+		std::vector<MockTreeWidgetItem*> groups_;
+		std::map<TreeWidgetItem*, std::vector<MockTreeWidgetItem*> > groupMembers_;
+		MockTreeWidget* root_;
 };
 
 }
diff --git a/Swiften/Roster/UnitTest/MockTreeWidgetItem.h b/Swiften/Roster/UnitTest/MockTreeWidgetItem.h
index 9b05d40..a40aca7 100644
--- a/Swiften/Roster/UnitTest/MockTreeWidgetItem.h
+++ b/Swiften/Roster/UnitTest/MockTreeWidgetItem.h
@@ -12,13 +12,16 @@ namespace Swift {
 class MockTreeWidgetItem : public TreeWidgetItem {
 	public:
 		virtual ~MockTreeWidgetItem() {};
-		virtual void setText(const String&) {};
+		virtual void setText(const String& text) {text_ = text;};
+		String getText() {return text_;};
 		virtual void setStatusText(const String&) {};
 		virtual void setAvatarPath(const String&) {};
 		virtual void setExpanded(bool) {};
 		virtual void setStatusShow(StatusShow::Type /*show*/) {};
 		virtual void show() {};
 		virtual void hide() {};
+	private:
+		String text_;
 };
 
 }