diff --git a/slixmpp/plugins/xep_0047/stream.py b/slixmpp/plugins/xep_0047/stream.py index eee6bd0e..3be894eb 100644 --- a/slixmpp/plugins/xep_0047/stream.py +++ b/slixmpp/plugins/xep_0047/stream.py @@ -31,7 +31,8 @@ class IBBytestream(object): self.recv_queue = asyncio.Queue() - def send(self, data): + @asyncio.coroutine + def send(self, data, timeout=None): if not self.stream_started or self.stream_out_closed: raise socket.error if len(data) > self.block_size: @@ -55,17 +56,22 @@ class IBBytestream(object): iq['ibb_data']['sid'] = self.sid iq['ibb_data']['seq'] = seq iq['ibb_data']['data'] = data - iq.send(callback=self._recv_ack) + yield from iq.send(timeout=timeout) return len(data) - def sendall(self, data): + @asyncio.coroutine + def sendall(self, data, timeout=None): sent_len = 0 while sent_len < len(data): - sent_len += self.send(data[sent_len:self.block_size]) + sent_len += yield from self.send(data[sent_len:self.block_size], timeout=timeout) - def _recv_ack(self, iq): - if iq['type'] == 'error': - self.close() + @asyncio.coroutine + def sendfile(self, file, timeout=None): + while True: + data = file.read(self.block_size) + if not data: + break + yield from self.send(data, timeout=timeout) def _recv_data(self, stanza): new_seq = stanza['ibb_data']['seq'] @@ -80,7 +86,7 @@ class IBBytestream(object): raise XMPPError('not-acceptable') self.recv_queue.put_nowait(data) - self.xmpp.event('ibb_stream_data', {'stream': self, 'data': data}) + self.xmpp.event('ibb_stream_data', self) if isinstance(stanza, Iq): stanza.reply().send() @@ -93,7 +99,7 @@ class IBBytestream(object): raise socket.error return self.recv_queue.get_nowait() - def close(self): + def close(self, timeout=None): iq = self.xmpp.Iq() iq['type'] = 'set' iq['to'] = self.peer_jid @@ -102,8 +108,9 @@ class IBBytestream(object): self.stream_out_closed = True def _close_stream(_): self.stream_in_closed = True - iq.send(callback=_close_stream) + future = iq.send(timeout=timeout, callback=_close_stream) self.xmpp.event('ibb_stream_end', self) + return future def _closed(self, iq): self.stream_in_closed = True diff --git a/tests/test_stream_xep_0047.py b/tests/test_stream_xep_0047.py index 2cc43823..ecba2445 100644 --- a/tests/test_stream_xep_0047.py +++ b/tests/test_stream_xep_0047.py @@ -1,3 +1,4 @@ +import asyncio import threading import time @@ -78,6 +79,7 @@ class TestInBandByteStreams(SlixTest): self.assertEqual(events, set(['ibb_stream_start', 'callback'])) + @asyncio.coroutine def testSendData(self): """Test sending data over an in-band bytestream.""" @@ -115,7 +117,7 @@ class TestInBandByteStreams(SlixTest): # Test sending data out - stream.send("Testing") + yield from stream.send("Testing") self.send("""