plugin XEP-0231: separated the requestData code in a method which can called independantly + some minor improvments

src/plugins/plugin_xep_0231.py

@@ -19,12 +19,16 @@
 
 from sat.core.i18n import _
 from sat.core.constants import Const as C
+from sat.core import exceptions
 from sat.core.log import getLogger
 log = getLogger(__name__)
 from sat.tools import xml_tools
 from wokkel import disco, iwokkel
 from zope.interface import implements
-from twisted.words.protocols.jabber.xmlstream import XMPPHandler
+from twisted.python import failure
+from twisted.words.protocols.jabber import xmlstream
+from twisted.words.protocols.jabber import jid
+from functools import partial
 import base64
 
 
@@ -49,27 +53,28 @@ class XEP_0231(object):
         self.host = host
         host.trigger.add("xhtml_post_treat", self.XHTMLTrigger)
 
-    def dumpData(self, client, data_elt, cid):
+    def dumpData(self, cache, data_elt, cid):
         """save file encoded in data_elt to cache
 
+        @param cache(memory.cache.Cache): cache to use to store the data
         @param data_elt(domish.Element): <data> as in XEP-0231
         @param cid(unicode): content-id
         @return(unicode): full path to dumped file
         """
         # FIXME: is it needed to use a separate thread?
         #        probably not with the little data expected with BoB
-        mime_type = data_elt.getAttribute('type','')
-
         try:
             max_age = int(data_elt['max-age'])
+            if max_age < 0:
+                raise ValueError
         except (KeyError, ValueError):
             log.warning(u'invalid max-age found')
             max_age = None
 
-        with client.cache.cacheData(
-            PLUGIN_INFO['import_name'],
+        with cache.cacheData(
+            PLUGIN_INFO[C.PI_IMPORT_NAME],
             cid,
-            mime_type,
+            data_elt.getAttribute('type'),
             max_age) as f:
 
             file_path = f.name
@@ -80,20 +85,43 @@ class XEP_0231(object):
     def getHandler(self, client):
         return XEP_0231_handler()
 
-    def _dataCb(self, iq_elt, client, img_elt, cid):
+    def _requestCb(self, iq_elt, cache, cid):
         for data_elt in iq_elt.elements(NS_BOB, u'data'):
             if data_elt.getAttribute('cid') == cid:
-                file_path = self.dumpData(client, data_elt, cid)
-                img_elt[u'src'] = u'file://{}'.format(file_path)
-                break
-        else:
-            log.warning(u"invalid data stanza received, requested cid was not found:\n{iq_elt}\nrequested cid: {cid}".format(
-                iq_elt = iq_elt,
-                cid = cid
-                ))
+                file_path = self.dumpData(cache, data_elt, cid)
+                return file_path
 
-    def _dataEb(self, iq_elt):
-        log.warning(u"Can't get requested data:\n{iq_elt}".format(iq_elt=iq_elt))
+        log.warning(u"invalid data stanza received, requested cid was not found:\n{iq_elt}\nrequested cid: {cid}".format(
+            iq_elt = iq_elt,
+            cid = cid
+            ))
+        raise failure.Failure(exceptions.DataError("missing data"))
+
+    def _requestEb(self, failure_):
+        """Log the error and continue errback chain"""
+        log.warning(u"Can't get requested data:\n{reason}".format(reason=failure_))
+        return failure_
+
+    def requestData(self, client, to_jid, cid, cache=None):
+        """Request data if we don't have it in cache
+
+        @param to_jid(jid.JID): jid to request the data to
+        @param cid(unicode): content id
+        @param cache(memory.cache.Cache, None): cache to use
+            client.cache will be used if None
+        """
+        if cache is None:
+            cache = client.cache
+        iq_elt = client.IQ('get')
+        iq_elt['to'] = to_jid.full()
+        data_elt = iq_elt.addElement((NS_BOB, 'data'))
+        data_elt['cid'] = cid
+        d = iq_elt.send()
+        d.addCallback(self._requestCb, cache, cid)
+        d.addErrback(self._requestEb)
+
+    def _setImgEltSrc(self, path, img_elt):
+        img_elt[u'src'] = u'file://{}'.format(path)
 
     def XHTMLTrigger(self, client, message_elt, body_elt, lang, treat_d):
         for img_elt in xml_tools.findAll(body_elt, C.NS_XHTML, u'img'):
@@ -109,23 +137,18 @@ class XEP_0231(object):
                     # image is not in cache, is it given locally?
                     for data_elt in message_elt.elements(NS_BOB, u'data'):
                         if data_elt.getAttribute('cid') == cid:
-                            file_path = self.dumpData(data_elt, cid)
+                            file_path = self.dumpData(client.cache, data_elt, cid)
                             img_elt[u'src'] = u'file://{}'.format(file_path)
                             break
                     else:
                         # cid not found locally, we need to request it
                         # so we use the deferred
-                        iq_elt = client.IQ('get')
-                        iq_elt['to'] = message_elt['from']
-                        data_elt = iq_elt.addElement((NS_BOB, 'data'))
-                        data_elt['cid'] = cid
-                        d = iq_elt.send()
-                        d.addCallback(self._dataCb, client, img_elt, cid)
-                        d.addErrback(self._dataEb)
+                        d = self.requestData(client, jid.JID(message_elt['from']), cid)
+                        d.addCallback(partial(self._setImgEltSrc, img_elt=img_elt))
                         treat_d.addCallback(lambda dummy: d)
 
 
-class XEP_0231_handler(XMPPHandler):
+class XEP_0231_handler(xmlstream.XMPPHandler):
     implements(iwokkel.IDisco)
 
     def getDiscoInfo(self, requestor, target, nodeIdentifier=''):