Fix detecting end of result set paging.

2 files changed, 52 insertions, 26 deletions

diff --git a/sleekxmpp/plugins/xep_0030/stanza/items.py b/sleekxmpp/plugins/xep_0030/stanza/items.py
index a1fb819c..f642a5b8 100644
--- a/sleekxmpp/plugins/xep_0030/stanza/items.py
+++ b/sleekxmpp/plugins/xep_0030/stanza/items.py
@@ -6,7 +6,7 @@
     See the file LICENSE for copying permission.
 """
 
-from sleekxmpp.xmlstream import ElementBase, ET
+from sleekxmpp.xmlstream import ElementBase, ET, register_stanza_plugin
 
 
 class DiscoItems(ElementBase):
@@ -78,13 +78,11 @@ class DiscoItems(ElementBase):
         """
         if (jid, node) not in self._items:
             self._items.add((jid, node))
-            item_xml = ET.Element('{%s}item' % self.namespace)
-            item_xml.attrib['jid'] = jid
-            if name:
-                item_xml.attrib['name'] = name
-            if node:
-                item_xml.attrib['node'] = node
-            self.xml.append(item_xml)
+            item = DiscoItem(parent=self)
+            item['jid'] = jid
+            item['node'] = node
+            item['name'] = name
+            self.iterables.append(item)
             return True
         return False
 
@@ -108,11 +106,9 @@ class DiscoItems(ElementBase):
     def get_items(self):
         """Return all items."""
         items = set()
-        for item_xml in self.findall('{%s}item' % self.namespace):
-            item = (item_xml.attrib['jid'],
-                    item_xml.attrib.get('node'),
-                    item_xml.attrib.get('name'))
-            items.add(item)
+        for item in self['substanzas']:
+            if isinstance(item, DiscoItem):
+                items.add((item['jid'], item['node'], item['name']))
         return items
 
     def set_items(self, items):
@@ -132,5 +128,23 @@ class DiscoItems(ElementBase):
     def del_items(self):
         """Remove all items."""
         self._items = set()
-        for item_xml in self.findall('{%s}item' % self.namespace):
-            self.xml.remove(item_xml)
+        for item in self['substanzas']:
+            if isinstance(item, DiscoItem):
+                self.xml.remove(item.xml)
+
+
+class DiscoItem(ElementBase):
+    name = 'item'
+    namespace = 'http://jabber.org/protocol/disco#items'
+    plugin_attrib = name
+    interfaces = set(('jid', 'node', 'name'))
+
+    def get_node(self):
+        return self._get_attr('node', None)
+
+    def get_name(self):
+        return self._get_attr('name', None)
+
+
+register_stanza_plugin(DiscoItems, DiscoItem, iterable=True)
+
diff --git a/sleekxmpp/plugins/xep_0059/rsm.py b/sleekxmpp/plugins/xep_0059/rsm.py
index 35908473..0e62bdde 100644
--- a/sleekxmpp/plugins/xep_0059/rsm.py
+++ b/sleekxmpp/plugins/xep_0059/rsm.py
@@ -13,6 +13,7 @@ from sleekxmpp import Iq
 from sleekxmpp.plugins.base import base_plugin
 from sleekxmpp.xmlstream import register_stanza_plugin
 from sleekxmpp.plugins.xep_0059 import Set
+from sleekxmpp.exceptions import XMPPError
 
 
 log = logging.getLogger(__name__)
@@ -70,19 +71,30 @@ class ResultIterator():
         elif self.start:
             self.query[self.interface]['rsm']['after'] = self.start
 
-        r = self.query.send(block=True)
-
-        if not r or not r[self.interface]['rsm']['first'] and \
-           not r[self.interface]['rsm']['last']:
+        try:
+            r = self.query.send(block=True)
+
+            if not r[self.interface]['rsm']['first'] and \
+               not r[self.interface]['rsm']['last']:
+                raise StopIteration
+
+            if r[self.interface]['rsm']['count'] and \
+               r[self.interface]['rsm']['first_index']:
+                count = int(r[self.interface]['rsm']['count'])
+                first = int(r[self.interface]['rsm']['first_index'])
+                num_items = len(r[self.interface]['substanzas'])
+                if first + num_items == count:
+                    raise StopIteration
+            
+            if self.reverse:
+                self.start = r[self.interface]['rsm']['first']
+            else:
+                self.start = r[self.interface]['rsm']['last']
+
+            return r
+        except XMPPError:
             raise StopIteration
 
-        if self.reverse:
-            self.start = r[self.interface]['rsm']['first']
-        else:
-            self.start = r[self.interface]['rsm']['last']
-
-        return r
-
 
 class xep_0059(base_plugin):