2 files changed, 28 insertions, 6 deletions
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_;
 };
 
 }