From ad5b61de50ca4fbcc445ceae8db568be6518e66e Mon Sep 17 00:00:00 2001
From: Lance Stout <lancestout@gmail.com>
Date: Mon, 30 Apr 2012 11:07:54 -0700
Subject: [PATCH] Add full support for initial payloads with adhoc commands,
 plus test.

---
 sleekxmpp/plugins/xep_0050/adhoc.py | 32 ++++++++++++++---
 tests/test_stream_xep_0050.py       | 54 +++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 5 deletions(-)

diff --git a/sleekxmpp/plugins/xep_0050/adhoc.py b/sleekxmpp/plugins/xep_0050/adhoc.py
index dc99e343..fcf51b1a 100644
--- a/sleekxmpp/plugins/xep_0050/adhoc.py
+++ b/sleekxmpp/plugins/xep_0050/adhoc.py
@@ -213,13 +213,21 @@ class XEP_0050(BasePlugin):
         name, handler = self.commands.get(key, ('Not found', None))
         if not handler:
             log.debug('Command not found: %s, %s', key, self.commands)
+
+        payload = []
+        for stanza in iq['command']['substanzas']:
+            payload.append(stanza)
+
+        interfaces = set([item.plugin_attrib for item in payload])
+        payload_classes = set([item.__class__ for item in payload])
+
         initial_session = {'id': sessionid,
                            'from': iq['from'],
                            'to': iq['to'],
                            'node': node,
-                           'payload': None,
-                           'interfaces': '',
-                           'payload_classes': None,
+                           'payload': payload,
+                           'interfaces': interfaces,
+                           'payload_classes': payload_classes,
                            'notes': None,
                            'has_next': False,
                            'allow_complete': False,
@@ -269,11 +277,19 @@ class XEP_0050(BasePlugin):
         sessionid = session['id']
 
         payload = session['payload']
+        if payload is None:
+            payload = []
         if not isinstance(payload, list):
             payload = [payload]
 
-        session['interfaces'] = [item.plugin_attrib for item in payload]
-        session['payload_classes'] = [item.__class__ for item in payload]
+        interfaces = session.get('interfaces', set())
+        payload_classes = session.get('payload_classes', set())
+
+        interfaces.update(set([item.plugin_attrib for item in payload]))
+        payload_classes.update(set([item.__class__ for item in payload]))
+
+        session['interfaces'] = interfaces
+        session['payload_classes'] = payload_classes
 
         self.sessions[sessionid] = session
 
@@ -486,6 +502,12 @@ class XEP_0050(BasePlugin):
         session['from'] = ifrom
         iq['command']['node'] = node
         iq['command']['action'] = 'execute'
+        if session['payload'] is not None:
+            payload = session['payload']
+            if not isinstance(payload, list):
+                payload = list(payload)
+            for stanza in payload:
+                iq['command'].append(stanza)
         sessionid = 'client:pending_' + iq['id']
         session['id'] = sessionid
         self.sessions[sessionid] = session
diff --git a/tests/test_stream_xep_0050.py b/tests/test_stream_xep_0050.py
index 1931349d..373bce64 100644
--- a/tests/test_stream_xep_0050.py
+++ b/tests/test_stream_xep_0050.py
@@ -1,4 +1,5 @@
 import time
+import logging
 import threading
 
 from sleekxmpp.test import *
@@ -17,6 +18,59 @@ class TestAdHocCommands(SleekTest):
     def tearDown(self):
         self.stream_close()
 
+    def testInitialPayloadCommand(self):
+        """Test a command with an initial payload."""
+
+        class TestPayload(ElementBase):
+            name = 'foo'
+            namespace = 'test'
+            interfaces = set(['bar'])
+            plugin_attrib = name
+
+        Command = self.xmpp['xep_0050'].stanza.Command
+        register_stanza_plugin(Command, TestPayload, iterable=True)
+
+        def handle_command(iq, session):
+            initial = session['payload']
+            logging.debug(initial)
+            new_payload = TestPayload()
+            if initial:
+                new_payload['bar'] = 'Received: %s' % initial[0]['bar'] 
+            else:
+                new_payload['bar'] = 'Failed'
+
+            logging.debug(initial)
+
+            session['payload'] = new_payload
+            session['next'] = None
+            session['has_next'] = False
+
+            return session
+
+        self.xmpp['xep_0050'].add_command('tester@localhost', 'foo',
+                                          'Do Foo', handle_command)
+
+        self.recv("""
+          <iq id="11" type="set" to="tester@localhost" from="foo@bar">
+            <command xmlns="http://jabber.org/protocol/commands"
+                     node="foo"
+                     action="execute">
+              <foo xmlns="test" bar="baz" />
+            </command>
+          </iq>
+        """)
+
+        self.send("""
+          <iq id="11" type="result" to="foo@bar">
+            <command xmlns="http://jabber.org/protocol/commands"
+                     node="foo"
+                     status="completed"
+                     sessionid="_sessionid_">
+              <foo xmlns="test" bar="Received: baz" />
+            </command>
+          </iq>
+        """)
+
     def testZeroStepCommand(self):
         """Test running a command with no steps."""