Updated, cleaned, and documented Iq stanza class. Also added unit tests.

sleekxmpp/stanza/iq.py

@@ -5,73 +5,170 @@
 
     See the file LICENSE for copying permission.
 """
-from .. xmlstream.stanzabase import StanzaBase
+
-from xml.etree import cElementTree as ET
+from sleekxmpp.stanza import Error
-from . error import Error
+from sleekxmpp.stanza.rootstanza import RootStanza
-from .. xmlstream.handler.waiter import Waiter
+from sleekxmpp.xmlstream import RESPONSE_TIMEOUT
-from .. xmlstream.matcher.id import MatcherId
+from sleekxmpp.xmlstream.stanzabase import StanzaBase, ET
-from . rootstanza import RootStanza
+from sleekxmpp.xmlstream.handler import Waiter
+from sleekxmpp.xmlstream.matcher import MatcherId
+
 
 class Iq(RootStanza):
-	interfaces = set(('type', 'to', 'from', 'id','query'))
-	types = set(('get', 'result', 'set', 'error'))
-	name = 'iq'
-	plugin_attrib = name
-	namespace = 'jabber:client'
 
-	def __init__(self, *args, **kwargs):
+    """
-		StanzaBase.__init__(self, *args, **kwargs)
+    XMPP <iq> stanzas, or info/query stanzas, are XMPP's method of
-		if self['id'] == '':
+    requesting and modifying information, similar to HTTP's GET and
-			if self.stream is not None:
+    POST methods.
-				self['id'] = self.stream.getNewId()
-			else:
-				self['id'] = '0'
 
-	def unhandled(self):
+    Each <iq> stanza must have an 'id' value which associates the
-		if self['type'] in ('get', 'set'):
+    stanza with the response stanza. XMPP entities must always
-			self.reply()
+    be given a response <iq> stanza with a type of 'result' after
-			self['error']['condition'] = 'feature-not-implemented'
+    sending a stanza of type 'get' or 'set'.
-			self['error']['text'] = 'No handlers registered for this request.'
-			self.send()
 
-	def setPayload(self, value):
+    Most uses cases for <iq> stanzas will involve adding a <query>
-		self.clear()
+    element whose namespace indicates the type of information
-		StanzaBase.setPayload(self, value)
+    desired. However, some custom XMPP applications use <iq> stanzas
-		return self
+    as a carrier stanza for an application-specific protocol instead.
 
-	def setQuery(self, value):
+    Example <iq> Stanzas:
-		query = self.xml.find("{%s}query" % value)
+        <iq to="user@example.com" type="get" id="314">
-		if query is None and value:
+          <query xmlns="http://jabber.org/protocol/disco#items" />
-			self.clear()
+        </iq>
-			query = ET.Element("{%s}query" % value)
-			self.xml.append(query)
-		return self
 
-	def getQuery(self):
+        <iq to="user@localhost" type="result" id="17">
-		for child in self.xml.getchildren():
+          <query xmlns='jabber:iq:roster'>
-			if child.tag.endswith('query'):
+            <item jid='otheruser@example.net'
-				ns =child.tag.split('}')[0]
+                  name='John Doe'
-				if '{' in ns:
+                  subscription='both'>
-					ns = ns[1:]
+              <group>Friends</group>
-				return ns
+            </item>
-		return ''
+          </query>
+        </iq>
 
-	def reply(self):
+    Stanza Interface:
-		self['type'] = 'result'
+        query -- The namespace of the <query> element if one exists.
-		StanzaBase.reply(self)
-		return self
 
-	def delQuery(self):
+    Methods:
-		for child in self.getchildren():
+        __init__   -- Overrides StanzaBase.__init__.
-			if child.tag.endswith('query'):
+        unhandled  -- Send error if there are no handlers.
-				self.xml.remove(child)
+        setPayload -- Overrides StanzaBase.setPayload.
-		return self
+        setQuery   -- Add or modify a <query> element.
+        getQuery   -- Return the namespace of the <query> element.
+        delQuery   -- Remove the <query> element.
+        reply      -- Overrides StanzaBase.reply
+        send       -- Overrides StanzaBase.send
+    """
 
-	def send(self, block=True, timeout=10):
+    namespace = 'jabber:client'
-		if block and self['type'] in ('get', 'set'):
+    name = 'iq'
-			waitfor = Waiter('IqWait_%s' % self['id'], MatcherId(self['id']))
+    interfaces = set(('type', 'to', 'from', 'id', 'query'))
-			self.stream.registerHandler(waitfor)
+    types = set(('get', 'result', 'set', 'error'))
-			StanzaBase.send(self)
+    plugin_attrib = name
-			return waitfor.wait(timeout)
+
-		else:
+    def __init__(self, *args, **kwargs):
-			return StanzaBase.send(self)
+        """
+        Initialize a new <iq> stanza with an 'id' value.
+
+        Overrides StanzaBase.__init__.
+        """
+        StanzaBase.__init__(self, *args, **kwargs)
+        if self['id'] == '':
+            if self.stream is not None:
+                self['id'] = self.stream.getNewId()
+            else:
+                self['id'] = '0'
+
+    def unhandled(self):
+        """
+        Send a feature-not-implemented error if the stanza is not handled.
+
+        Overrides StanzaBase.unhandled.
+        """
+        if self['type'] in ('get', 'set'):
+            self.reply()
+            self['error']['condition'] = 'feature-not-implemented'
+            self['error']['text'] = 'No handlers registered for this request.'
+            self.send()
+
+    def setPayload(self, value):
+        """
+        Set the XML contents of the <iq> stanza.
+
+        Arguments:
+            value -- An XML object to use as the <iq> stanza's contents
+        """
+        self.clear()
+        StanzaBase.setPayload(self, value)
+        return self
+
+    def setQuery(self, value):
+        """
+        Add or modify a <query> element.
+
+        Query elements are differentiated by their namespace.
+
+        Arguments:
+            value -- The namespace of the <query> element.
+        """
+        query = self.xml.find("{%s}query" % value)
+        if query is None and value:
+            self.clear()
+            query = ET.Element("{%s}query" % value)
+            self.xml.append(query)
+        return self
+
+    def getQuery(self):
+        """Return the namespace of the <query> element."""
+        for child in self.xml.getchildren():
+            if child.tag.endswith('query'):
+                ns = child.tag.split('}')[0]
+                if '{' in ns:
+                    ns = ns[1:]
+                return ns
+        return ''
+
+    def delQuery(self):
+        """Remove the <query> element."""
+        for child in self.xml.getchildren():
+            if child.tag.endswith('query'):
+                self.xml.remove(child)
+        return self
+
+    def reply(self):
+        """
+        Send a reply <iq> stanza.
+
+        Overrides StanzaBase.reply
+
+        Sets the 'type' to 'result' in addition to the default
+        StanzaBase.reply behavior.
+        """
+        self['type'] = 'result'
+        StanzaBase.reply(self)
+        return self
+
+    def send(self, block=True, timeout=RESPONSE_TIMEOUT):
+        """
+        Send an <iq> stanza over the XML stream.
+
+        The send call can optionally block until a response is received or
+        a timeout occurs. Be aware that using blocking in non-threaded event
+        handlers can drastically impact performance.
+
+        Overrides StanzaBase.send
+
+        Arguments:
+            block   -- Specify if the send call will block until a response
+                       is received, or a timeout occurs. Defaults to True.
+            timeout -- The length of time (in seconds) to wait for a response
+                       before exiting the send call if blocking is used.
+                       Defaults to sleekxmpp.xmlstream.RESPONSE_TIMEOUT
+        """
+        if block and self['type'] in ('get', 'set'):
+            waitfor = Waiter('IqWait_%s' % self['id'], MatcherId(self['id']))
+            self.stream.registerHandler(waitfor)
+            StanzaBase.send(self)
+            return waitfor.wait(timeout)
+        else:
+            return StanzaBase.send(self)

tests/test_iqstanzas.py

@@ -0,0 +1,93 @@
+from sleektest import *
+from sleekxmpp.xmlstream.stanzabase import ET
+
+
+class TestIqStanzas(SleekTest):
+
+    def setUp(self):
+        """Start XML stream for testing."""
+        self.streamStart()
+
+    def tearDown(self):
+        """Shutdown the XML stream after testing."""
+        self.streamClose()
+
+    def testSetup(self):
+        """Test initializing default Iq values."""
+        iq = self.Iq()
+        self.checkIq(iq, """
+          <iq id="0" />
+        """)
+    
+    def testPayload(self):
+        """Test setting Iq stanza payload."""
+        iq = self.Iq()
+        iq.setPayload(ET.Element('{test}tester'))
+        self.checkIq(iq, """
+          <iq id="0">
+            <tester xmlns="test" />
+          </iq>
+        """, use_values=False)
+
+
+    def testUnhandled(self):
+        """Test behavior for Iq.unhandled."""
+        self.streamRecv("""
+          <iq id="test" type="get">
+            <query xmlns="test" />
+           </iq>
+        """)
+
+	iq = self.Iq()
+        iq['id'] = 'test'
+	iq['error']['condition'] = 'feature-not-implemented'
+	iq['error']['text'] = 'No handlers registered for this request.'
+
+        self.streamSendIq(iq, """
+          <iq id="test" type="error">
+            <error type="cancel">
+              <feature-not-implemented xmlns="urn:ietf:params:xml:ns:xmpp-stanzas" />
+              <text xmlns="urn:ietf:params:xml:ns:xmpp-stanzas">
+                No handlers registered for this request.
+              </text>
+            </error>
+          </iq>
+        """)
+
+    def testQuery(self):
+        """Test modifying query element of Iq stanzas."""
+        iq = self.Iq()
+
+        iq['query'] = 'query_ns'
+        self.checkIq(iq, """
+          <iq id="0">
+            <query xmlns="query_ns" />
+          </iq>
+        """)
+
+        iq['query'] = 'query_ns2'
+        self.checkIq(iq, """
+          <iq id="0">
+            <query xmlns="query_ns2" />
+          </iq>
+        """)
+
+	self.failUnless(iq['query'] == 'query_ns2', "Query namespace doesn't match")
+
+	del iq['query']
+	self.checkIq(iq, """
+          <iq id="0" />
+        """)
+
+    def testReply(self):
+        """Test setting proper result type in Iq replies."""
+        iq = self.Iq()
+	iq['to'] = 'user@localhost'
+	iq['type'] = 'get'
+        iq.reply()
+
+	self.checkIq(iq, """
+          <iq id="0" type="result" />
+        """)
+
+suite = unittest.TestLoader().loadTestsFromTestCase(TestIqStanzas)