1 files changed, 24 insertions, 5 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_;
 };
 
 }