4 files changed, 46 insertions, 7 deletions
diff --git a/Swiften/Roster/UnitTest/MockTreeWidget.h b/Swiften/Roster/UnitTest/MockTreeWidget.h
index b766aeb..97dd796 100644
--- a/Swiften/Roster/UnitTest/MockTreeWidget.h
+++ b/Swiften/Roster/UnitTest/MockTreeWidget.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.
  */
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_;
diff --git a/Swiften/Roster/UnitTest/MockTreeWidgetItem.cpp b/Swiften/Roster/UnitTest/MockTreeWidgetItem.cpp
new file mode 100644
index 0000000..4238ead
--- /dev/null
+++ b/Swiften/Roster/UnitTest/MockTreeWidgetItem.cpp
@@ -0,0 +1,19 @@
+/*
+ * Copyright (c) 2010 Kevin Smith
+ * Licensed under the GNU General Public License v3.
+ * See Documentation/Licenses/GPLv3.txt for more information.
+ */
+#include "Swiften/Roster/UnitTest/MockTreeWidgetItem.h"
+
+#include "Swiften/Roster/UnitTest/MockTreeWidgetFactory.h"
+
+namespace Swift {
+
+MockTreeWidgetItem::~MockTreeWidgetItem() {
+	factory_->removeItem(this);
+}
+
+}
+
+
+
diff --git a/Swiften/Roster/UnitTest/MockTreeWidgetItem.h b/Swiften/Roster/UnitTest/MockTreeWidgetItem.h
index 1ec1441..08b9f1a 100644
--- a/Swiften/Roster/UnitTest/MockTreeWidgetItem.h
+++ b/Swiften/Roster/UnitTest/MockTreeWidgetItem.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.
  */
@@ -14,10 +14,11 @@
 #include <boost/shared_ptr.hpp>
 
 namespace Swift {
-
+	class MockTreeWidgetFactory;
 class MockTreeWidgetItem : public TreeWidgetItem {
 	public:
-		virtual ~MockTreeWidgetItem() {};
+		MockTreeWidgetItem(MockTreeWidgetFactory* factory) {factory_ = factory;};
+		virtual ~MockTreeWidgetItem();
 		virtual void setText(const String& text) {text_ = text;};
 		String getText() {return text_;};
 		virtual void setStatusText(const String&) {};
@@ -28,6 +29,7 @@ class MockTreeWidgetItem : public TreeWidgetItem {
 		virtual void hide() {};
 	private:
 		String text_;
+		MockTreeWidgetFactory* factory_;
 };
 
 }