1 files changed, 21 insertions, 3 deletions
diff --git a/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h b/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h
index 935423c..d94c859 100644
--- a/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h
+++ b/Swiften/Roster/UnitTest/MockTreeWidgetFactory.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2010 Remko Tronçon
+ * Copyright (c) 2010 Kevin Smith
  * Licensed under the GNU General Public License v3.
  * See Documentation/Licenses/GPLv3.txt for more information.
  */
@@ -27,12 +27,12 @@ class MockTreeWidgetFactory : public TreeWidgetFactory {
 			return root_;
 		};
 		virtual TreeWidgetItem* createTreeWidgetItem(TreeWidgetItem* group) {
-			MockTreeWidgetItem* entry = new MockTreeWidgetItem();
+			MockTreeWidgetItem* entry = new MockTreeWidgetItem(this);
 			groupMembers_[group].push_back(entry);
 			return entry;
 		};
 		virtual TreeWidgetItem* createTreeWidgetItem(TreeWidget*) {
-			MockTreeWidgetItem* group = new MockTreeWidgetItem();
+			MockTreeWidgetItem* group = new MockTreeWidgetItem(this);
 			groups_.push_back(group);
 			return group;
 		};
@@ -43,6 +43,24 @@ class MockTreeWidgetFactory : public TreeWidgetFactory {
 			}
 			return groupNames;
 		};
+
+		typedef std::map<TreeWidgetItem*, std::vector<MockTreeWidgetItem*> > itemMap;
+
+		virtual std::vector<MockTreeWidgetItem*> getGroupMembers(const String& group) {
+			for (itemMap::iterator it = groupMembers_.begin(); it != groupMembers_.end(); it++) {
+				if (((MockTreeWidgetItem*)(it->first))->getText() == group) {
+					return it->second;
+				}
+			}
+			return std::vector<MockTreeWidgetItem*>();
+		};
+
+		virtual void removeItem(MockTreeWidgetItem* item) {
+			foreach (TreeWidgetItem* groupItem, groups_) {
+				std::vector<MockTreeWidgetItem*>& members = groupMembers_[groupItem];
+				members.erase(std::remove(members.begin(), members.end(), item), members.end());
+			}
+		};
 	private:
 		std::vector<MockTreeWidgetItem*> groups_;
 		std::map<TreeWidgetItem*, std::vector<MockTreeWidgetItem*> > groupMembers_;