3 files changed, 40 insertions, 7 deletions
diff --git a/Swiften/Roster/GroupRosterItem.cpp b/Swiften/Roster/GroupRosterItem.cpp
index 05530ec..b55ce17 100644
--- a/Swiften/Roster/GroupRosterItem.cpp
+++ b/Swiften/Roster/GroupRosterItem.cpp
@@ -76,7 +76,7 @@ bool GroupRosterItem::itemLessThan(const RosterItem* left, const RosterItem* rig
 	const ContactRosterItem* leftContact = dynamic_cast<const ContactRosterItem*>(left);
 	const ContactRosterItem* rightContact = dynamic_cast<const ContactRosterItem*>(right);
 	if (leftContact) {
-		if (rightContact) {
+		if (!rightContact) {
 			return false;
 		}
 		StatusShow::Type leftType = leftContact->getSimplifiedStatusShow();
diff --git a/Swiften/Roster/Roster.cpp b/Swiften/Roster/Roster.cpp
index c25fd41..9226a63 100644
--- a/Swiften/Roster/Roster.cpp
+++ b/Swiften/Roster/Roster.cpp
@@ -75,15 +75,15 @@ void Roster::addContact(const JID& jid, const String& name, const String& groupN
 	filterContact(item, group);
 }
 
+struct JIDEqualsTo {
+	JIDEqualsTo(const JID& jid) : jid(jid) {}
+	bool operator()(ContactRosterItem* i) const { return jid == i->getJID(); }
+	JID jid;
+};
 
 void Roster::removeContact(const JID& jid) {
 	std::vector<ContactRosterItem*> items = itemMap_[jid.toBare()];
-	std::vector<ContactRosterItem*>::iterator it = items.begin();
-	while (it != items.end()) {
-		if (jid == (*it)->getJID()) {
-			it = items.erase(it);
-		}
-	}
+	items.erase(std::remove_if(items.begin(), items.end(), JIDEqualsTo(jid)), items.end());
 	if (items.size() == 0) {
 		itemMap_.erase(jid.toBare());
 	}
diff --git a/Swiften/Roster/UnitTest/RosterTest.cpp b/Swiften/Roster/UnitTest/RosterTest.cpp
index f9dff33..2697747 100644
--- a/Swiften/Roster/UnitTest/RosterTest.cpp
+++ b/Swiften/Roster/UnitTest/RosterTest.cpp
@@ -17,6 +17,9 @@ class RosterTest : public CppUnit::TestFixture
 {
 		CPPUNIT_TEST_SUITE(RosterTest);
 		CPPUNIT_TEST(testGetGroup);
+		CPPUNIT_TEST(testRemoveContact);
+		CPPUNIT_TEST(testRemoveSecondContact);
+		CPPUNIT_TEST(testRemoveSecondContactSameBare);
 		CPPUNIT_TEST_SUITE_END();
 
 	private:
@@ -51,6 +54,36 @@ class RosterTest : public CppUnit::TestFixture
 
 		}
 
+		void testRemoveContact() {
+			roster_->addContact(jid1_, "Bert", "group1");
+			CPPUNIT_ASSERT_EQUAL(String("Bert"), ((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren()[0]->getDisplayName());
+
+			roster_->removeContact(jid1_);
+			CPPUNIT_ASSERT_EQUAL(0, (int)((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren().size());
+		}
+
+		void testRemoveSecondContact() {
+			roster_->addContact(jid1_, "Bert", "group1");
+			roster_->addContact(jid2_, "Cookie", "group1");
+			CPPUNIT_ASSERT_EQUAL(String("Cookie"), ((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren()[1]->getDisplayName());
+
+			roster_->removeContact(jid2_);
+			CPPUNIT_ASSERT_EQUAL(1, (int)((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren().size());
+			CPPUNIT_ASSERT_EQUAL(String("Bert"), ((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren()[0]->getDisplayName());
+		}
+
+		void testRemoveSecondContactSameBare() {
+			JID jid4a("a@b/c");
+			JID jid4b("a@b/d");
+			roster_->addContact(jid4a, "Bert", "group1");
+			roster_->addContact(jid4b, "Cookie", "group1");
+			CPPUNIT_ASSERT_EQUAL(String("Cookie"), ((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren()[1]->getDisplayName());
+
+			roster_->removeContact(jid4b);
+			CPPUNIT_ASSERT_EQUAL(1, (int)((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren().size());
+			CPPUNIT_ASSERT_EQUAL(String("Bert"), ((GroupRosterItem*)roster_->getRoot()->getChildren()[0])->getChildren()[0]->getDisplayName());
+		}
+
 };
 CPPUNIT_TEST_SUITE_REGISTRATION(RosterTest);