First pass at integrating the new roster manager.

9 files changed, 514 insertions, 137 deletions

diff --git a/sleekxmpp/basexmpp.py b/sleekxmpp/basexmpp.py
index cd7d2510..fb718fba 100644
--- a/sleekxmpp/basexmpp.py
+++ b/sleekxmpp/basexmpp.py
@@ -15,6 +15,7 @@ import logging
 import sleekxmpp
 from sleekxmpp import plugins
 
+from sleekxmpp.roster import MultiRoster
 from sleekxmpp.stanza import Message, Presence, Iq, Error
 from sleekxmpp.stanza.roster import Roster
 from sleekxmpp.stanza.nick import Nick
@@ -78,7 +79,7 @@ class BaseXMPP(XMLStream):
        send_presence_subscribe -- Send a subscription request.
     """
 
-    def __init__(self, default_ns='jabber:client'):
+    def __init__(self, jid='', default_ns='jabber:client'):
         """
         Adapt an XML stream for use with XMPP.
 
@@ -107,10 +108,13 @@ class BaseXMPP(XMLStream):
         self.default_ns = default_ns
         self.stream_ns = 'http://etherx.jabber.org/streams'
 
-        self.boundjid = JID("")
+        self.boundjid = JID(jid)
 
         self.plugin = {}
+        self.rosters = MultiRoster(self)
+        self.rosters.add(self.boundjid.bare)
         self.roster = {}
+
         self.is_component = False
         self.auto_authorize = True
         self.auto_subscribe = True
@@ -127,10 +131,20 @@ class BaseXMPP(XMLStream):
                      MatchXPath("{%s}presence" % self.default_ns),
                      self._handle_presence))
 
-        self.add_event_handler('presence_subscribe',
-                               self._handle_subscribe)
         self.add_event_handler('disconnected',
                                self._handle_disconnected)
+        self.add_event_handler('presence_available', self._handle_available)
+        self.add_event_handler('presence_dnd', self._handle_available)
+        self.add_event_handler('presence_xa', self._handle_available)
+        self.add_event_handler('presence_chat', self._handle_available)
+        self.add_event_handler('presence_away', self._handle_available)
+        self.add_event_handler('presence_unavailable', self._handle_unavailable)
+        self.add_event_handler('presence_subscribe', self._handle_subscribe)
+        self.add_event_handler('presence_subscribed', self._handle_subscribed)
+        self.add_event_handler('presence_unsubscribe', self._handle_unsubscribe)
+        self.add_event_handler('presence_unsubscribed', self._handle_unsubscribed)
+        self.add_event_handler('presence_probe', self._handle_probe)
+        self.add_event_handler('roster_subscription_request', self._handle_new_subscription)
 
         # Set up the XML stream with XMPP's root stanzas.
         self.registerStanza(Message)
@@ -522,12 +536,49 @@ class BaseXMPP(XMLStream):
         """Process incoming message stanzas."""
         self.event('message', msg)
 
+    def _handle_available(self, presence):
+        self.rosters[presence['to'].bare][presence['from'].bare].handle_available(presence)
+
+    def _handle_unavailable(self, presence):
+        self.rosters[presence['to'].bare][presence['from'].bare].handle_unavailable(presence)
+
+    def _handle_new_subscription(self, stanza):
+        roster = self.rosters[stanza['to'].bare]
+        item = self.rosters[stanza['to'].bare][stanza['from'].bare]
+        if item['whitelisted']:
+            item.authorize()
+        elif roster.auto_authorize:
+            item.authorize()
+            if roster.auto_subscribe:
+                item.subscribe()
+        elif roster.auto_authorize == False:
+            item.unauthorize()
+
+    def _handle_removed_subscription(self, presence):
+        self.rosters[presence['to'].bare][presence['from'].bare].unauthorize()
+
+    def _handle_subscribe(self, stanza):
+        self.rosters[stanza['to'].bare][stanza['from'].bare].handle_subscribe(stanza)
+
+    def _handle_subscribed(self, stanza):
+        self.rosters[stanza['to'].bare][stanza['from'].bare].handle_subscribed(stanza)
+
+    def _handle_unsubscribe(self, stanza):
+        self.rosters[stanza['to'].bare][stanza['from'].bare].handle_unsubscribe(stanza)
+
+    def _handle_unsubscribed(self, stanza):
+        self.rosters[stanza['to'].bare][stanza['from'].bare].handle_unsubscribed(stanza)
+
+    def _handle_probe(self, stanza):
+        self.rosteritems[stanza['to'].bare][stanza['from'].bare].handle_probe(stanza)
+
     def _handle_presence(self, presence):
         """
         Process incoming presence stanzas.
 
         Update the roster with presence information.
         """
+        logging.debug(presence['type'])
         self.event("presence_%s" % presence['type'], presence)
 
         # Check for changes in subscription state.
@@ -538,97 +589,7 @@ class BaseXMPP(XMLStream):
         elif not presence['type'] in ('available', 'unavailable') and \
              not presence['type'] in presence.showtypes:
             return
-
-        # Strip the information from the stanza.
-        jid = presence['from'].bare
-        resource = presence['from'].resource
-        show = presence['type']
-        status = presence['status']
-        priority = presence['priority']
-
-        was_offline = False
-        got_online = False
-        old_roster = self.roster.get(jid, {}).get(resource, {})
-
-        # Create a new roster entry if needed.
-        if not jid in self.roster:
-            self.roster[jid] = {'groups': [],
-                                'name': '',
-                                'subscription': 'none',
-                                'presence': {},
-                                'in_roster': False}
-
-        # Alias to simplify some references.
-        connections = self.roster[jid]['presence']
-
-        # Determine if the user has just come online.
-        if not resource in connections:
-            if show == 'available' or show in presence.showtypes:
-                got_online = True
-            was_offline = True
-            connections[resource] = {}
-
-        if connections[resource].get('show', 'unavailable') == 'unavailable':
-            was_offline = True
-
-        # Update the roster's state for this JID's resource.
-        connections[resource] = {'show': show,
-                                'status': status,
-                                'priority': priority}
-
-        name = self.roster[jid].get('name', '')
-
-        # Remove unneeded state information after a resource
-        # disconnects. Determine if this was the last connection
-        # for the JID.
-        if show == 'unavailable':
-            log.debug("%s %s got offline" % (jid, resource))
-            del connections[resource]
-
-            if not connections and not self.roster[jid]['in_roster']:
-                del self.roster[jid]
-            if not was_offline:
-                self.event("got_offline", presence)
-            else:
-                return False
-
-        name = '(%s) ' % name if name else ''
-
-        # Presence state has changed.
         self.event("changed_status", presence)
-        if got_online:
-            self.event("got_online", presence)
-        log.debug("STATUS: %s%s/%s[%s]: %s" % (name, jid, resource,
-                                                   show, status))
-
-    def _handle_subscribe(self, presence):
-        """
-        Automatically managage subscription requests.
-
-        Subscription behavior is controlled by the settings
-        self.auto_authorize and self.auto_subscribe.
-
-        auto_auth  auto_sub   Result:
-        True       True       Create bi-directional subsriptions.
-        True       False      Create only directed subscriptions.
-        False      *          Decline all subscriptions.
-        None       *          Disable automatic handling and use
-                              a custom handler.
-        """
-        presence.reply()
-        presence['to'] = presence['to'].bare
-
-        # We are using trinary logic, so conditions have to be
-        # more explicit than usual.
-        if self.auto_authorize == True:
-            presence['type'] = 'subscribed'
-            presence.send()
-            if self.auto_subscribe:
-                presence['type'] = 'subscribe'
-                presence.send()
-        elif self.auto_authorize == False:
-            presence['type'] = 'unsubscribed'
-            presence.send()
 
 # Restore the old, lowercased name for backwards compatibility.
 basexmpp = BaseXMPP
diff --git a/sleekxmpp/clientxmpp.py b/sleekxmpp/clientxmpp.py
index 1c600812..8b8f05f6 100644
--- a/sleekxmpp/clientxmpp.py
+++ b/sleekxmpp/clientxmpp.py
@@ -66,7 +66,7 @@ class ClientXMPP(BaseXMPP):
                                 when calling register_plugins.
             escape_quotes    -- Deprecated.
         """
-        BaseXMPP.__init__(self, 'jabber:client')
+        BaseXMPP.__init__(self, jid, 'jabber:client')
 
         # To comply with PEP8, method names now use underscores.
         # Deprecated method names are re-mapped for backwards compatibility.
@@ -75,7 +75,6 @@ class ClientXMPP(BaseXMPP):
         self.getRoster = self.get_roster
         self.registerFeature = self.register_feature
 
-        self.set_jid(jid)
         self.password = password
         self.escape_quotes = escape_quotes
         self.plugin_config = plugin_config
@@ -421,13 +420,13 @@ class ClientXMPP(BaseXMPP):
         """
         if iq['type'] == 'set' or (iq['type'] == 'result' and request):
             for jid in iq['roster']['items']:
-                if not jid in self.roster:
-                    self.roster[jid] = {'groups': [],
-                                        'name': '',
-                                        'subscription': 'none',
-                                        'presence': {},
-                                        'in_roster': True}
-                self.roster[jid].update(iq['roster']['items'][jid])
+                item = iq['roster']['items'][jid]
+                roster = self.rosters[iq['to'].bare]
+                roster[jid]['name'] = item['name']
+                roster[jid]['groups'] = item['groups']
+                roster[jid]['from'] = item['subscription'] in ['from', 'both']
+                roster[jid]['to'] = item['subscription'] in ['to', 'both']
+                roster[jid]['pending_out'] = (item['ask'] == 'subscribe')
 
         self.event("roster_update", iq)
         if iq['type'] == 'set':
diff --git a/sleekxmpp/componentxmpp.py b/sleekxmpp/componentxmpp.py
index ae58c5f2..0963c502 100644
--- a/sleekxmpp/componentxmpp.py
+++ b/sleekxmpp/componentxmpp.py
@@ -58,7 +58,7 @@ class ComponentXMPP(BaseXMPP):
             default_ns = 'jabber:client'
         else:
             default_ns = 'jabber:component:accept'
-        BaseXMPP.__init__(self, default_ns)
+        BaseXMPP.__init__(self, jid, default_ns)
 
         self.auto_authorize = None
         self.stream_header = "<stream:stream %s %s to='%s'>" % (
@@ -68,8 +68,8 @@ class ComponentXMPP(BaseXMPP):
         self.stream_footer = "</stream:stream>"
         self.server_host = host
         self.server_port = port
-        self.set_jid(jid)
         self.secret = secret
+
         self.plugin_config = plugin_config
         self.plugin_whitelist = plugin_whitelist
         self.is_component = True
diff --git a/sleekxmpp/roster.py b/sleekxmpp/roster.py
new file mode 100644
index 00000000..32fd5c0f
--- /dev/null
+++ b/sleekxmpp/roster.py
@@ -0,0 +1,374 @@
+import logging
+
+
+class MultiRoster(object):
+
+    def __init__(self, xmpp, datastore=None):
+        self.xmpp = xmpp
+        self.datastore = datastore
+        self._rosters = {}
+
+    def __getitem__(self, key):
+        if key not in self._rosters:
+            self.add(key)
+        return self._rosters[key]
+
+    def keys(self):
+        return self._rosters.keys()
+
+    def __iter__(self):
+        return self._rosters.__iter__()
+
+    def add(self, node):
+        if node not in self._rosters:
+            self._rosters[node] = Roster(self.xmpp, node, self.datastore)
+
+class Roster(object):
+
+    def __init__(self, xmpp, jid, datastore=None):
+        self.xmpp = xmpp
+        self.jid = jid
+        self.datastore = datastore
+        self.auto_authorize = True
+        self.auto_subscribe = True
+        self._jids = {}
+
+    def __getitem__(self, key):
+        if key not in self._jids:
+            self.add(key, save=True)
+        return self._jids[key]
+
+    def keys(self):
+        return self._jids.keys()
+
+    def __iter__(self):
+        return self._jids.__iter__()
+
+    def add(self, jid, name='', groups=None, afrom=False, ato=False,
+            pending_in=False, pending_out=False, whitelisted=False,
+            save=False):
+        state = {'name': name,
+                 'groups': groups or [],
+                 'from': afrom,
+                 'to': ato,
+                 'pending_in': pending_in,
+                 'pending_out': pending_out,
+                 'whitelisted': whitelisted,
+                 'subscription': 'none'}
+        self._jids[jid] = RosterItem(self.xmpp, jid, self.jid,
+                                     state=state, datastore=self.datastore)
+        if save:
+            self._jids[jid].save()
+
+    def subscribe(self, jid):
+        self._jids[jid].subscribe()
+
+    def unsubscribe(self, jid):
+        self._jids[jid].unsubscribe()
+
+    def remove(self, jid):
+        self._jids[jid].remove()
+        if not self.xmpp.is_component:
+            self.update(jid, subscription='remove')
+
+    def update(self, jid, name=None, subscription=None, groups=[]):
+        self._jids[jid]['name'] = name
+        self._jids[jid]['groups'] = group
+        self._jids[jid].save()
+
+        if not self.xmpp.is_component:
+            iq = self.Iq()
+            iq['type'] = 'set'
+            iq['roster']['items'] = {jid: {'name': name,
+                                           'subscription': subscription,
+                                           'groups': groups}}
+            response = iq.send()
+            return response and response['type'] == 'result'
+
+    def presence(self, jid, resource=None):
+        if resource is None:
+            return self._jids[jid].resources
+
+        default_presence = {'status': '',
+                            'priority': 0,
+                            'show': ''}
+        return self._jids[jid].resources.get(resource,
+                                             default_presence)
+
+
+class RosterItem(object):
+
+    def __init__(self, xmpp, jid, owner=None,
+                 state=None, datastore=None):
+        self.xmpp = xmpp
+        self.jid = jid
+        self.owner = owner or self.xmpp.jid
+        self.last_status = None
+        self.resources = {}
+        self.datastore = datastore
+
+        self._state = state or {
+                'from': False,
+                'to': False,
+                'pending_in': False,
+                'pending_out': False,
+                'whitelisted': False,
+                'subscription': 'none',
+                'name': '',
+                'groups': []}
+        self._datastore_state = {}
+        self.load()
+
+    def load(self):
+        if self.datastore:
+            item = self.datastore.load(self.owner, self.jid,
+                                       self._datastore_state)
+            if item:
+                self['name'] = item['name']
+                self['groups'] = item['groups']
+                self['from'] = item['from']
+                self['to'] = item['to']
+                self['whitelisted'] = item['whitelisted']
+                self['pending_out'] = item['pending_out']
+                self['pending_in'] = item['pending_in']
+                self['subscription'] = self._subscription()
+            return self._state
+        return None
+
+    def save(self):
+        if self.datastore:
+            self.datastore.save(self.owner, self.jid,
+                                self._state, self._datastore_state)
+
+    def __getitem__(self, key):
+        if key in self._state:
+            if key == 'subscription':
+                return self._subscription()
+            return self._state[key]
+        else:
+            raise KeyError
+
+    def __setitem__(self, key, value):
+        print "%s: %s" % (key, value)
+        if key in self._state:
+            if key in ['name', 'subscription', 'groups']:
+                self._state[key] = value
+            else:
+                value = str(value).lower()
+                self._state[key] = value in ('true', '1', 'on', 'yes')
+        else:
+            raise KeyError
+
+    def _subscription(self):
+        if self['to'] and self['from']:
+            return 'both'
+        elif self['from']:
+            return 'from'
+        elif self['to']:
+            return 'to'
+        else:
+            return 'none'
+
+    def remove(self):
+        "Remove the jids subscription, inform it if it is subscribed, and unwhitelist it"
+        if self['to']:
+            p = self.xmpp.Presence()
+            p['to'] = self.jid
+            p['type'] = ['unsubscribe']
+            if self.xmpp.is_component:
+                p['from'] = self.owner
+            p.send()
+            self['to'] = False
+        self['whitelisted'] = False
+        self.save()
+
+    def subscribe(self):
+        p = self.xmpp.Presence()
+        p['to'] = self.jid
+        p['type'] = 'subscribe'
+        if self.xmpp.is_component:
+            p['from'] = self.owner
+        self['pending_out'] = True
+        self.save()
+        p.send()
+
+    def authorize(self):
+        self['from'] = True
+        self['pending_in'] = False
+        self.save()
+        self._subscribed()
+        self.send_last_presence()
+
+    def unauthorize(self):
+        self['from'] = False
+        self['pending_in'] = False
+        self.save()
+        self._unsubscribed()
+        p = self.xmpp.Presence()
+        p['to'] = self.jid
+        p['type'] = 'unavailable'
+        if self.xmpp.is_component:
+            p['from'] = self.owner
+        p.send()
+
+    def _subscribed(self):
+        p = self.xmpp.Presence()
+        p['to'] = self.jid
+        p['type'] = 'subscribed'
+        if self.xmpp.is_component:
+            p['from'] = self.owner
+        p.send()
+
+    def unsubscribe(self):
+        p = self.xmpp.Presence()
+        p['to'] = self.jid
+        p['type'] = 'unsubscribe'
+        if self.xmpp.is_component:
+            p['from'] = self.owner
+        self.save()
+        p.send()
+
+    def _unsubscribed(self):
+        p = self.xmpp.Presence()
+        p['to'] = self.jid
+        p['type'] = 'unsubscribed'
+        if self.xmpp.is_component:
+            p['from'] = self.owner
+        p.send()
+
+    def send_presence(self, ptype='available', status=None):
+        p = self.xmpp.Presence()
+        p['to'] = self.jid
+        p['type'] = ptype
+        p['status'] = status
+        if self.xmpp.is_component:
+            p['from'] = self.owner
+        self.last_status = p
+        p.send()
+
+    def send_last_presence(self):
+        if self.last_status is None:
+            self.send_presence()
+        else:
+            self.last_status.send()
+
+    def handle_available(self, presence):
+        resource = presence['from'].resource
+        data = {'status': presence['status'],
+                'show': presence['show'],
+                'priority': presence['priority']}
+        if not self.resources:
+            self.xmpp.event('got_online', presence)
+        if resource not in self.resources:
+            self.resources[resource] = {}
+        self.resources[resource].update(data)
+
+    def handle_unavailable(self, presence):
+        resource = presence['from'].resource
+        if not self.resources:
+            return
+        if resource in self.resources:
+            del self.resources[resource]
+        if not self.resources:
+            self.xmpp.event('got_offline', presence)
+
+    def handle_subscribe(self, presence):
+        """
+        +------------------------------------------------------------------+
+        |  EXISTING STATE          |  DELIVER?  |  NEW STATE               |
+        +------------------------------------------------------------------+
+        |  "None"                  |  yes       |  "None + Pending In"     |
+        |  "None + Pending Out"    |  yes       |  "None + Pending Out/In" |
+        |  "None + Pending In"     |  no        |  no state change         |
+        |  "None + Pending Out/In" |  no        |  no state change         |
+        |  "To"                    |  yes       |  "To + Pending In"       |
+        |  "To + Pending In"       |  no        |  no state change         |
+        |  "From"                  |  no *      |  no state change         |
+        |  "From + Pending Out"    |  no *      |  no state change         |
+        |  "Both"                  |  no *      |  no state change         |
+        +------------------------------------------------------------------+
+        """
+        if not self['from'] and not self['pending_in']:
+            self['pending_in'] = True
+            self.xmpp.event('roster_subscription_request', presence)
+        elif self['from']:
+            self._subscribed()
+        self.save()
+
+    def handle_subscribed(self, presence):
+        """
+        +------------------------------------------------------------------+
+        |  EXISTING STATE          |  DELIVER?  |  NEW STATE               |
+        +------------------------------------------------------------------+
+        |  "None"                  |  no        |  no state change         |
+        |  "None + Pending Out"    |  yes       |  "To"                    |
+        |  "None + Pending In"     |  no        |  no state change         |
+        |  "None + Pending Out/In" |  yes       |  "To + Pending In"       |
+        |  "To"                    |  no        |  no state change         |
+        |  "To + Pending In"       |  no        |  no state change         |
+        |  "From"                  |  no        |  no state change         |
+        |  "From + Pending Out"    |  yes       |  "Both"                  |
+        |  "Both"                  |  no        |  no state change         |
+        +------------------------------------------------------------------+
+        """
+        if not self['to'] and self['pending_out']:
+            self['pending_out'] = False
+            self['to'] = True
+            self.xmpp.event('roster_subscription_authorized', presence)
+        self.save()
+
+    def handle_unsubscribe(self, presence):
+        """
+        +------------------------------------------------------------------+
+        |  EXISTING STATE          |  DELIVER?  |  NEW STATE               |
+        +------------------------------------------------------------------+
+        |  "None"                  |  no        |  no state change         |
+        |  "None + Pending Out"    |  no        |  no state change         |
+        |  "None + Pending In"     |  yes *     |  "None"                  |
+        |  "None + Pending Out/In" |  yes *     |  "None + Pending Out"    |
+        |  "To"                    |  no        |  no state change         |
+        |  "To + Pending In"       |  yes *     |  "To"                    |
+        |  "From"                  |  yes *     |  "None"                  |
+        |  "From + Pending Out"    |  yes *     |  "None + Pending Out     |
+        |  "Both"                  |  yes *     |  "To"                    |
+        +------------------------------------------------------------------+
+        """
+        if not self['from']  and self['pending_in']:
+            self['pending_in'] = False
+            self._unsubscribed()
+        elif self['from']:
+            self['from'] = False
+            self._unsubscribed()
+            self.xmpp.event('roster_subscription_remove', presence)
+        self.save()
+
+    def handle_unsubscribed(self, presence):
+        """
+        +------------------------------------------------------------------+
+        |  EXISTING STATE          |  DELIVER?  |  NEW STATE               |
+        +------------------------------------------------------------------+
+        |  "None"                  |  no        |  no state change         |
+        |  "None + Pending Out"    |  yes       |  "None"                  |
+        |  "None + Pending In"     |  no        |  no state change         |
+        |  "None + Pending Out/In" |  yes       |  "None + Pending In"     |
+        |  "To"                    |  yes       |  "None"                  |
+        |  "To + Pending In"       |  yes       |  "None + Pending In"     |
+        |  "From"                  |  no        |  no state change         |
+        |  "From + Pending Out"    |  yes       |  "From"                  |
+        |  "Both"                  |  yes       |  "From"                  |
+        +------------------------------------------------------------------
+        """
+        if not self['to'] and self['pending_out']:
+            self['pending_out'] = False
+        elif self['to'] and not self['pending_out']:
+            self['to'] = False
+            self.xmpp.event('roster_subscription_removed', presence)
+        self.save()
+
+    def handle_probe(self, presence):
+        if self['to']:
+            self.send_last_presence()
+        if self['pending_out']:
+            self.subscribe()
+        if not self['to']:
+            self._unsubscribed()
diff --git a/sleekxmpp/stanza/roster.py b/sleekxmpp/stanza/roster.py
index 8f154a22..57ea62ba 100644
--- a/sleekxmpp/stanza/roster.py
+++ b/sleekxmpp/stanza/roster.py
@@ -106,6 +106,8 @@ class Roster(ElementBase):
                 item = {}
                 item['name'] = itemxml.get('name', '')
                 item['subscription'] = itemxml.get('subscription', '')
+                item['ask'] = itemxml.get('ask', '')
+                item['approved'] = itemxml.get('approved', '')
                 item['groups'] = []
                 groupsxml = itemxml.findall('{jabber:iq:roster}group')
                 if groupsxml is not None:
diff --git a/sleekxmpp/test/sleektest.py b/sleekxmpp/test/sleektest.py
index f8b4b546..c481d73b 100644
--- a/sleekxmpp/test/sleektest.py
+++ b/sleekxmpp/test/sleektest.py
@@ -137,6 +137,33 @@ class SleekTest(unittest.TestCase):
             self.assertEqual(str(jid), string,
                     "String does not match: %s" % str(jid))
 
+    def check_roster(self, owner, jid, name=None, subscription=None,
+                     afrom=None, ato=None, pending_out=None, pending_in=None,
+                     groups=None):
+        roster = self.xmpp.rosters[owner][jid]
+        print roster._state
+        if name is not None:
+            self.assertEqual(roster['name'], name,
+                    "Incorrect name value: %s" % roster['name'])
+        if subscription is not None:
+            self.assertEqual(roster['subscription'], subscription,
+                    "Incorrect subscription: %s" % roster['subscription'])
+        if afrom is not None:
+            self.assertEqual(roster['from'], afrom,
+                    "Incorrect from state: %s" % roster['from'])
+        if ato is not None:
+            self.assertEqual(roster['to'], ato,
+                    "Incorrect to state: %s" % roster['to'])
+        if pending_out is not None:
+            self.assertEqual(roster['pending_out'], pending_out,
+                    "Incorrect pending_out state: %s" % roster['pending_out'])
+        if pending_in is not None:
+            self.assertEqual(roster['pending_in'], pending_out,
+                    "Incorrect pending_in state: %s" % roster['pending_in'])
+        if groups is not None:
+            self.assertEqual(roster['groups'], groups,
+                    "Incorrect groups: %s" % roster['groups'])
+
     # ------------------------------------------------------------------
     # Methods for comparing stanza objects to XML strings
 
diff --git a/tests/test_stanza_roster.py b/tests/test_stanza_roster.py
index cd3e607c..8ec2d32b 100644
--- a/tests/test_stanza_roster.py
+++ b/tests/test_stanza_roster.py
@@ -48,10 +48,14 @@ class TestRosterStanzas(SleekTest):
             'user@example.com': {
                 'name': 'User',
                 'subscription': 'both',
+                'ask': '',
+                'approved': '',
                 'groups': ['Friends', 'Coworkers']},
             'otheruser@example.com': {
                 'name': 'Other User',
                 'subscription': 'both',
+                'ask': '',
+                'approved': '',
                 'groups': []}}
         debug = "Roster items don't match after retrieval."
         debug += "\nReturned: %s" % str(iq['roster']['items'])
diff --git a/tests/test_stream_presence.py b/tests/test_stream_presence.py
index 1d5caa98..8ca1c0e9 100644
--- a/tests/test_stream_presence.py
+++ b/tests/test_stream_presence.py
@@ -30,7 +30,9 @@ class TestStreamPresence(SleekTest):
         self.xmpp.add_event_handler('presence_unavailable', unavailable)
 
         self.recv("""
-          <presence type="unavailable" from="otheruser@localhost" />
+          <presence type="unavailable"
+                    from="otheruser@localhost"
+                    to="tester@localhost"/>
         """)
 
         # Give event queue time to process.
@@ -68,12 +70,14 @@ class TestStreamPresence(SleekTest):
 
         # Contact comes online.
         self.recv("""
-          <presence from="otheruser@localhost/foobar" />
+          <presence from="otheruser@localhost/foobar"
+                    to="tester@localhost" />
         """)
 
         # Contact goes offline, should trigger got_offline.
         self.recv("""
           <presence from="otheruser@localhost/foobar"
+                    to="tester@localhost"
                     type="unavailable" />
         """)
 
@@ -99,7 +103,8 @@ class TestStreamPresence(SleekTest):
         self.xmpp.add_event_handler('got_online', got_online)
 
         self.recv("""
-          <presence from="user@localhost" />
+          <presence from="user@localhost"
+                    to="tester@localhost" />
         """)
 
         # Give event queue time to process.
@@ -136,15 +141,23 @@ class TestStreamPresence(SleekTest):
         self.xmpp.auto_subscribe = True
 
         self.recv("""
-          <presence from="user@localhost" type="subscribe" />
+          <presence from="user@localhost"
+                    to="tester@localhost"
+                    type="subscribe" />
+        """)
+
+        self.send("""
+          <presence to="user@localhost"
+                    type="subscribed" />
         """)
 
         self.send("""
-          <presence to="user@localhost" type="subscribed" />
+          <presence to="user@localhost" />
         """)
 
         self.send("""
-          <presence to="user@localhost" type="subscribe" />
+          <presence to="user@localhost"
+                    type="subscribe" />
         """)
 
         expected = set(('presence_subscribe', 'changed_subscription'))
@@ -170,14 +183,17 @@ class TestStreamPresence(SleekTest):
                                     presence_subscribe)
 
         # With this setting we should reject all subscriptions.
-        self.xmpp.auto_authorize = False
+        self.xmpp.rosters['tester@localhost'].auto_authorize = False
 
         self.recv("""
-          <presence from="user@localhost" type="subscribe" />
+          <presence from="user@localhost"
+                    to="tester@localhost"
+                    type="subscribe" />
         """)
 
         self.send("""
-          <presence to="user@localhost" type="unsubscribed" />
+          <presence to="user@localhost"
+                    type="unsubscribed" />
         """)
 
         expected = set(('presence_subscribe', 'changed_subscription'))
diff --git a/tests/test_stream_roster.py b/tests/test_stream_roster.py
index 165a8bc9..8ffd86aa 100644
--- a/tests/test_stream_roster.py
+++ b/tests/test_stream_roster.py
@@ -13,8 +13,7 @@ class TestStreamRoster(SleekTest):
 
     def testGetRoster(self):
         """Test handling roster requests."""
-        self.stream_start(mode='client')
-        self.failUnless(self.xmpp.roster == {}, "Initial roster not empty.")
+        self.stream_start(mode='client', jid='tester@localhost')
 
         # Since get_roster blocks, we need to run it in a thread.
         t = threading.Thread(name='get_roster', target=self.xmpp.get_roster)
@@ -26,11 +25,12 @@ class TestStreamRoster(SleekTest):
           </iq>
         """)
         self.recv("""
-          <iq type="result" id="1">
+          <iq to='tester@localhost' type="result" id="1">
             <query xmlns="jabber:iq:roster">
               <item jid="user@localhost"
                     name="User"
-                    subscription="both">
+                    subscription="from"
+                    ask="subscribe">
                 <group>Friends</group>
                 <group>Examples</group>
               </item>
@@ -41,21 +41,20 @@ class TestStreamRoster(SleekTest):
         # Wait for get_roster to return.
         t.join()
 
-        roster = {'user@localhost': {'name': 'User',
-                                     'subscription': 'both',
-                                     'groups': ['Friends', 'Examples'],
-                                     'presence': {},
-                                     'in_roster': True}}
-        self.failUnless(self.xmpp.roster == roster,
-                "Unexpected roster values: %s" % self.xmpp.roster)
+        print self.xmpp.rosters['tester@localhost']['user@localhost']._state
+        self.check_roster('tester@localhost', 'user@localhost',
+                          name='User',
+                          subscription='from',
+                          afrom=True,
+                          pending_out=True,
+                          groups=['Friends', 'Examples'])
 
     def testRosterSet(self):
         """Test handling pushed roster updates."""
-        self.stream_start(mode='client')
-        self.failUnless(self.xmpp.roster == {}, "Initial roster not empty.")
+        self.stream_start(mode='client', jid='tester@localhost')
 
         self.recv("""
-          <iq type="set" id="1">
+          <iq to='tester@localhost' type="set" id="1">
             <query xmlns="jabber:iq:roster">
               <item jid="user@localhost"
                     name="User"
@@ -72,15 +71,10 @@ class TestStreamRoster(SleekTest):
           </iq>
         """)
 
-        roster = {'user@localhost': {'name': 'User',
-                                     'subscription': 'both',
-                                     'groups': ['Friends', 'Examples'],
-                                     'presence': {},
-                                     'in_roster': True}}
-        self.failUnless(self.xmpp.roster == roster,
-                "Unexpected roster values: %s" % self.xmpp.roster)
-
-
+        self.check_roster('tester@localhost', 'user@localhost',
+                          name='User',
+                          subscription='both',
+                          groups=['Friends', 'Examples'])
 
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestStreamRoster)