├── .coveragerc ├── .gitignore ├── .travis.yml ├── AUTHORS ├── CONTRIBUITORS ├── LICENSE ├── Makefile ├── README.md ├── example ├── subscribe.py └── torstomp ├── setup.py ├── tests ├── test_main.py └── test_recv_frame.py ├── torstomp ├── __init__.py ├── errors.py ├── frame.py ├── protocol.py └── subscription.py └── tox.ini /.coveragerc: -------------------------------------------------------------------------------- 1 | [run] 2 | omit = 3 | test_*.py 4 | branch = True 5 | source = 6 | torstomp 7 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | .coverage 3 | *.egg-info/ 4 | *.egg 5 | /htmlcov 6 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | cache: 3 | directories: 4 | - $HOME/.cache/pip 5 | sudo: false 6 | python: 7 | - "2.7" 8 | - "3.3" 9 | - "3.4" 10 | - "3.5" 11 | - "pypy" 12 | 13 | install: 14 | - make setup 15 | - pip install coveralls 16 | 17 | script: make test 18 | 19 | after_success: 20 | coveralls -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- 1 | Wilson Pinto Júnior 2 | -------------------------------------------------------------------------------- /CONTRIBUITORS: -------------------------------------------------------------------------------- 1 | Wilson Pinto Júnior 2 | Klaus Laube 3 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | Copyright (c) 2016 Wilson Pinto Júnior 3 | 4 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 5 | 6 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 7 | 8 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 9 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: setup clean test test_unit flake8 autopep8 upload 2 | 3 | setup: 4 | @pip install -Ue .\[tests\] 5 | 6 | clean: 7 | find . -name "*.pyc" -exec rm '{}' ';' 8 | 9 | unit test_unit test: 10 | @coverage run --branch `which nosetests` -v --with-yanc -s tests/ 11 | @$(MAKE) coverage 12 | @$(MAKE) static 13 | 14 | focus: 15 | @coverage run --branch `which nosetests` -vv --with-yanc --logging-level=WARNING --with-focus -i -s tests/ 16 | 17 | coverage: 18 | @coverage report -m --fail-under=80 19 | 20 | coverage_html: 21 | @coverage html 22 | @open htmlcov/index.html 23 | 24 | flake8 static: 25 | flake8 torstomp/ 26 | flake8 tests/ 27 | 28 | autopep8: 29 | autopep8 -r -i torstomp/ 30 | autopep8 -r -i tests/ 31 | 32 | upload: 33 | python ./setup.py sdist upload -r pypi 34 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![Build Status](https://travis-ci.org/wpjunior/torstomp.png?branch=master)](https://travis-ci.org/wpjunior/torstomp) 2 | [![Coverage Status](https://coveralls.io/repos/github/wpjunior/torstomp/badge.svg?branch=master)](https://coveralls.io/github/wpjunior/torstomp?branch=master) 3 | [![PyPI version](https://badge.fury.io/py/torstomp.svg)](https://badge.fury.io/py/torstomp) 4 | 5 | # Torstomp 6 | Simple tornado stomp 1.1 client. 7 | 8 | ## Install 9 | 10 | with pip: 11 | 12 | ```bash 13 | pip install torstomp 14 | ``` 15 | ## Usage 16 | ```python 17 | # -*- coding: utf-8 -*- 18 | 19 | from tornado import gen 20 | from tornado.ioloop import IOLoop 21 | from torstomp import TorStomp 22 | 23 | 24 | @gen.coroutine 25 | def main(): 26 | client = TorStomp('localhost', 61613, connect_headers={ 27 | 'heart-beat': '1000,1000' 28 | }, on_error=report_error) 29 | client.subscribe('/queue/channel', callback=on_message) 30 | 31 | yield client.connect() 32 | 33 | client.send('/queue/channel', body=u'Thanks', headers={}) 34 | 35 | 36 | def on_message(frame, message): 37 | print('on_message:', message) 38 | 39 | 40 | def report_error(error): 41 | print('report_error:', error) 42 | 43 | 44 | if __name__ == '__main__': 45 | main() 46 | IOLoop.current().start() 47 | ``` 48 | 49 | ## Development 50 | 51 | With empty virtualenv for this project, run this command: 52 | ```bash 53 | make setup 54 | ``` 55 | 56 | and run all tests =) 57 | ```bash 58 | make test 59 | ``` 60 | 61 | ## Contributing 62 | Fork, patch, test, and send a pull request. 63 | -------------------------------------------------------------------------------- /example/subscribe.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import logging 4 | 5 | from tornado import gen 6 | from tornado.ioloop import IOLoop 7 | 8 | from torstomp import TorStomp 9 | 10 | def on_message(frame, message): 11 | print(message) 12 | 13 | logging.basicConfig( 14 | format="%(asctime)s - %(filename)s:%(lineno)d - " 15 | "%(levelname)s - %(message)s", 16 | level='DEBUG') 17 | 18 | 19 | def report_error(error): 20 | print('report_error', error) 21 | 22 | @gen.coroutine 23 | def run(): 24 | client = TorStomp('localhost', 61613, connect_headers={ 25 | 'heart-beat': '1000,1000' 26 | }, on_error=report_error) 27 | client.subscribe('/queue/corumba', callback=on_message) 28 | 29 | yield client.connect() 30 | 31 | client.send('/queue/corumba', body=u'Wilson Júnior', headers={}) 32 | 33 | 34 | def main(): 35 | run() 36 | 37 | main() 38 | IOLoop.current().start() 39 | -------------------------------------------------------------------------------- /example/torstomp: -------------------------------------------------------------------------------- 1 | ../torstomp -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | 3 | from setuptools import find_packages, setup 4 | 5 | version = '0.1.12' 6 | 7 | setup( 8 | name='torstomp', 9 | version=version, 10 | description='Simple Stomp 1.1 client for tornado applications', 11 | long_description='', 12 | classifiers=[], 13 | keywords='stomp', 14 | author='Wilson Júnior', 15 | author_email='wilsonpjunior@gmail.com', 16 | url='https://github.com/wpjunior/torstomp.git', 17 | license='MIT', 18 | include_package_data=True, 19 | packages=find_packages(exclude=["tests", "tests.*"]), 20 | platforms=['any'], 21 | install_requires=[ 22 | 'six', 23 | 'tornado', 24 | ], 25 | extras_require={ 26 | 'tests': [ 27 | 'mock', 28 | 'nose', 29 | 'coverage', 30 | "yanc", 31 | "nose_focus", 32 | "flake8", 33 | ] 34 | } 35 | ) 36 | -------------------------------------------------------------------------------- /tests/test_main.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from torstomp import TorStomp 4 | from torstomp.subscription import Subscription 5 | from torstomp.frame import Frame 6 | 7 | from tornado.testing import AsyncTestCase, gen_test 8 | from tornado import gen 9 | from tornado.iostream import StreamClosedError 10 | 11 | from mock import MagicMock 12 | 13 | 14 | class TestTorStomp(AsyncTestCase): 15 | 16 | def setUp(self): 17 | super(TestTorStomp, self).setUp() 18 | self.stomp = TorStomp() 19 | 20 | def test_accept_version_header(self): 21 | self.assertEqual(self.stomp._connect_headers['accept-version'], '1.1') 22 | 23 | @gen_test 24 | def test_connect_write_subscriptions(self): 25 | callback = MagicMock() 26 | 27 | self.stomp._build_io_stream = MagicMock() 28 | 29 | io_stream = MagicMock() 30 | 31 | # mock write operation 32 | write_future = gen.Future() 33 | write_future.set_result(None) 34 | io_stream.write.return_value = write_future 35 | 36 | # mock connect operation 37 | connect_future = gen.Future() 38 | connect_future.set_result(None) 39 | io_stream.connect.return_value = connect_future 40 | 41 | self.stomp._build_io_stream.return_value = io_stream 42 | 43 | self.stomp.subscribe('/topic/test1', ack='client', extra_headers={ 44 | 'my-header': 'my-value' 45 | }, callback=callback) 46 | 47 | yield self.stomp.connect() 48 | 49 | write_calls = self.stomp.stream.write.call_args_list 50 | 51 | self.assertEqual( 52 | write_calls[0][0][0], 53 | b'CONNECT\naccept-version:1.1\n\n\x00' 54 | ) 55 | 56 | self.assertEqual( 57 | write_calls[1][0][0], 58 | b'SUBSCRIBE\n' 59 | b'ack:client\ndestination:/topic/test1\n' 60 | b'id:1\nmy-header:my-value\n\n' 61 | b'\x00' 62 | ) 63 | 64 | def test_subscribe_create_single_subscription(self): 65 | callback = MagicMock() 66 | 67 | self.stomp.stream = MagicMock() 68 | self.stomp.subscribe('/topic/test', ack='client', extra_headers={ 69 | 'my-header': 'my-value' 70 | }, callback=callback) 71 | 72 | subscription = self.stomp._subscriptions['1'] 73 | 74 | self.assertIsInstance(subscription, Subscription) 75 | self.assertEqual(subscription.destination, '/topic/test') 76 | self.assertEqual(subscription.id, 1) 77 | self.assertEqual(subscription.ack, 'client') 78 | self.assertEqual(subscription.extra_headers, {'my-header': 'my-value'}) 79 | self.assertEqual(subscription.callback, callback) 80 | 81 | def test_subscribe_create_multiple_subscriptions(self): 82 | callback1 = MagicMock() 83 | callback2 = MagicMock() 84 | 85 | self.stomp.stream = MagicMock() 86 | self.stomp.subscribe('/topic/test1', ack='client', extra_headers={ 87 | 'my-header': 'my-value' 88 | }, callback=callback1) 89 | 90 | self.stomp.subscribe('/topic/test2', callback=callback2) 91 | 92 | subscription = self.stomp._subscriptions['1'] 93 | 94 | self.assertIsInstance(subscription, Subscription) 95 | self.assertEqual(subscription.destination, '/topic/test1') 96 | self.assertEqual(subscription.id, 1) 97 | self.assertEqual(subscription.ack, 'client') 98 | self.assertEqual(subscription.extra_headers, {'my-header': 'my-value'}) 99 | self.assertEqual(subscription.callback, callback1) 100 | 101 | subscription = self.stomp._subscriptions['2'] 102 | 103 | self.assertIsInstance(subscription, Subscription) 104 | self.assertEqual(subscription.destination, '/topic/test2') 105 | self.assertEqual(subscription.id, 2) 106 | self.assertEqual(subscription.ack, 'auto') 107 | self.assertEqual(subscription.extra_headers, {}) 108 | self.assertEqual(subscription.callback, callback2) 109 | 110 | def test_subscribe_when_connected_write_in_stream(self): 111 | callback = MagicMock() 112 | 113 | self.stomp.stream = MagicMock() 114 | self.stomp.connected = True # fake connected 115 | self.stomp.subscribe('/topic/test', ack='client', extra_headers={ 116 | 'my-header': 'my-value' 117 | }, callback=callback) 118 | 119 | self.assertEqual(self.stomp.stream.write.call_count, 1) 120 | self.assertEqual( 121 | self.stomp.stream.write.call_args[0][0], 122 | b'SUBSCRIBE\n' 123 | b'ack:client\n' 124 | b'destination:/topic/test\n' 125 | b'id:1\n' 126 | b'my-header:my-value\n\n\x00') 127 | 128 | def test_subscribe_when_not_connected_write_in_stream(self): 129 | callback = MagicMock() 130 | 131 | self.stomp.stream = MagicMock() 132 | self.stomp.connected = False 133 | self.stomp.subscribe('/topic/test', ack='client', extra_headers={ 134 | 'my-header': 'my-value' 135 | }, callback=callback) 136 | 137 | self.assertEqual(self.stomp.stream.write.call_count, 0) 138 | 139 | def test_send_write_in_stream(self): 140 | self.stomp.stream = MagicMock() 141 | self.stomp.send('/topic/test', headers={ 142 | 'my-header': 'my-value' 143 | }, body='{}') 144 | 145 | self.assertEqual(self.stomp.stream.write.call_count, 1) 146 | self.assertEqual( 147 | self.stomp.stream.write.call_args[0][0], 148 | b'SEND\n' 149 | b'content-length:2\n' 150 | b'destination:/topic/test\n' 151 | b'my-header:my-value\n\n' 152 | b'{}\x00') 153 | 154 | def test_send_utf8_write_in_stream(self): 155 | self.stomp.stream = MagicMock() 156 | self.stomp.send('/topic/test', headers={ 157 | 'my-header': 'my-value' 158 | }, body=u'Wilson Júnior') 159 | 160 | self.assertEqual(self.stomp.stream.write.call_count, 1) 161 | self.assertEqual( 162 | self.stomp.stream.write.call_args[0][0], 163 | b'SEND\ncontent-length:14\ndestination:/topic/test\n' 164 | b'my-header:my-value\n\nWilson J\xc3\xbanior\x00') 165 | 166 | def test_jms_compatible_send_write_in_stream(self): 167 | self.stomp.stream = MagicMock() 168 | self.stomp.send('/topic/test', headers={ 169 | 'my-header': 'my-value', 170 | }, body='{}', send_content_length=False) 171 | 172 | self.assertEqual(self.stomp.stream.write.call_count, 1) 173 | self.assertEqual( 174 | self.stomp.stream.write.call_args[0][0], 175 | b'SEND\n' 176 | b'destination:/topic/test\n' 177 | b'my-header:my-value\n\n' 178 | b'{}\x00') 179 | 180 | def test_set_heart_beat_integration(self): 181 | self.stomp._set_heart_beat = MagicMock() 182 | self.stomp._on_data( 183 | b'CONNECTED\n' 184 | b'heart-beat:100,100\n\n' 185 | b'{}\x00') 186 | 187 | self.assertEqual(self.stomp._set_heart_beat.call_args[0][0], 100) 188 | 189 | def test_do_heart_beat(self): 190 | self.stomp.stream = MagicMock() 191 | self.stomp._schedule_heart_beat = MagicMock() 192 | self.stomp._do_heart_beat() 193 | 194 | self.assertEqual(self.stomp.stream.write.call_count, 1) 195 | self.assertEqual(self.stomp.stream.write.call_args[0][0], b'\n') 196 | 197 | self.assertEqual(self.stomp._schedule_heart_beat.call_count, 1) 198 | 199 | def test_do_heart_beat_with_closed_stream(self): 200 | self.stomp.stream = MagicMock() 201 | self.stomp.stream.write.side_effect = StreamClosedError() 202 | self.stomp._schedule_heart_beat = MagicMock() 203 | self.stomp._do_heart_beat() 204 | 205 | self.assertEqual(self.stomp.stream.write.call_count, 1) 206 | self.assertEqual(self.stomp.stream.write.call_args[0][0], b'\n') 207 | 208 | self.assertEqual(self.stomp._schedule_heart_beat.call_count, 1) 209 | 210 | def test_subscription_called(self): 211 | callback = MagicMock() 212 | 213 | self.stomp.stream = MagicMock() 214 | self.stomp.subscribe('/topic/test', ack='client', extra_headers={ 215 | 'my-header': 'my-value' 216 | }, callback=callback) 217 | 218 | self.stomp._on_data( 219 | b'MESSAGE\n' 220 | b'subscription:1\n' 221 | b'message-id:007\n' 222 | b'destination:/topic/test\n' 223 | b'\n' 224 | b'blah\x00') 225 | 226 | self.assertEqual(callback.call_count, 1) 227 | 228 | frame = callback.call_args[0][0] 229 | self.assertIsInstance(frame, Frame) 230 | self.assertEqual(frame.headers['message-id'], '007') 231 | self.assertEqual(frame.headers['subscription'], '1') 232 | self.assertEqual(callback.call_args[0][1], 'blah') 233 | 234 | def test_subscription_called_with_double_line_breaks_on_body(self): 235 | callback = MagicMock() 236 | 237 | self.stomp.stream = MagicMock() 238 | self.stomp.subscribe('/topic/test', ack='client', extra_headers={ 239 | 'my-header': 'my-value' 240 | }, callback=callback) 241 | 242 | self.stomp._on_data( 243 | b'MESSAGE\n' 244 | b'subscription:1\n' 245 | b'message-id:007\n' 246 | b'destination:/topic/test\n' 247 | b'\n' 248 | b'blahh-line-a\n\nblahh-line-b\n\nblahh-line-c\x00') 249 | 250 | self.assertEqual(callback.call_count, 1) 251 | 252 | frame = callback.call_args[0][0] 253 | self.assertIsInstance(frame, Frame) 254 | self.assertEqual(frame.headers['message-id'], '007') 255 | self.assertEqual(frame.headers['subscription'], '1') 256 | self.assertEqual(callback.call_args[0][1], 'blahh-line-a\n\nblahh-line-b\n\nblahh-line-c') 257 | 258 | def test_on_error_called(self): 259 | self.stomp._on_error = MagicMock() 260 | self.stomp._on_data( 261 | b'ERROR\n' 262 | b'message:Invalid error, blah, blah, blah\n' 263 | b'\n' 264 | b'Detail Error: blah, blah, blah\x00') 265 | 266 | self.assertEqual(self.stomp._on_error.call_count, 1) 267 | 268 | error = self.stomp._on_error.call_args[0][0] 269 | self.assertEqual(error.args[0], 'Invalid error, blah, blah, blah') 270 | self.assertEqual(error.detail, 'Detail Error: blah, blah, blah') 271 | 272 | def test_on_unhandled_frame(self): 273 | self.stomp._received_unhandled_frame = MagicMock() 274 | 275 | self.stomp._on_data( 276 | b'FIGHT\n' 277 | b'teste:1\n' 278 | b'\n' 279 | b'ok\x00') 280 | 281 | self.assertEqual(self.stomp._received_unhandled_frame.call_count, 1) 282 | 283 | frame = self.stomp._received_unhandled_frame.call_args[0][0] 284 | self.assertEqual(frame.command, 'FIGHT') 285 | self.assertEqual(frame.headers, {'teste': '1'}) 286 | self.assertEqual(frame.body, 'ok') 287 | 288 | def test_ack(self): 289 | self.stomp.stream = MagicMock() 290 | 291 | frame = Frame('MESSAGE', { 292 | 'subscription': '123', 293 | 'message-id': '321' 294 | }, 'blah') 295 | 296 | self.stomp.ack(frame) 297 | self.assertEqual(self.stomp.stream.write.call_count, 1) 298 | self.assertEqual( 299 | self.stomp.stream.write.call_args[0][0], 300 | b'ACK\n' 301 | b'message-id:321\n' 302 | b'subscription:123\n\n' 303 | b'\x00') 304 | 305 | def test_ack_with_unicode_headers(self): 306 | self.stomp.stream = MagicMock() 307 | 308 | frame = Frame('MESSAGE', { 309 | 'subscription': u'123', 310 | 'message-id': u'321' 311 | }, u'blah') 312 | 313 | self.stomp.ack(frame) 314 | self.assertEqual(self.stomp.stream.write.call_count, 1) 315 | buf = self.stomp.stream.write.call_args[0][0] 316 | self.assertIsInstance(buf, bytes) 317 | self.assertEqual( 318 | buf, 319 | b'ACK\n' 320 | b'message-id:321\n' 321 | b'subscription:123\n\n' 322 | b'\x00') 323 | 324 | def test_nack(self): 325 | self.stomp.stream = MagicMock() 326 | 327 | frame = Frame('MESSAGE', { 328 | 'subscription': '123', 329 | 'message-id': '321' 330 | }, 'blah') 331 | 332 | self.stomp.nack(frame) 333 | self.assertEqual(self.stomp.stream.write.call_count, 1) 334 | buf = self.stomp.stream.write.call_args[0][0] 335 | self.assertIsInstance(buf, bytes) 336 | self.assertEqual( 337 | buf, 338 | b'NACK\n' 339 | b'message-id:321\n' 340 | b'subscription:123\n\n' 341 | b'\x00') 342 | 343 | def test_unsubscribe(self): 344 | self.stomp.stream = MagicMock() 345 | 346 | self.stomp.subscribe("/queue/dummyqueue") 347 | 348 | last_subscribe_id = str(self.stomp._last_subscribe_id) 349 | subscription = self.stomp._subscriptions[last_subscribe_id] 350 | 351 | self.stomp.unsubscribe(subscription) 352 | 353 | self.assertFalse(last_subscribe_id in self.stomp._subscriptions.keys()) 354 | self.assertEqual(self.stomp.stream.write.call_count, 1) 355 | self.assertEqual( 356 | self.stomp.stream.write.call_args[0][0], 357 | b'UNSUBSCRIBE\n' 358 | b'destination:/queue/dummyqueue\n' 359 | b'id:1\n\n' 360 | b'\x00') 361 | -------------------------------------------------------------------------------- /tests/test_recv_frame.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | from unittest import TestCase 3 | import six 4 | 5 | from torstomp.protocol import StompProtocol 6 | 7 | from mock import MagicMock 8 | 9 | 10 | class TestRecvFrame(TestCase): 11 | 12 | def setUp(self): 13 | self.protocol = StompProtocol() 14 | 15 | def test_decode(self): 16 | self.assertEqual( 17 | self.protocol._decode(u'éĂ'), 18 | u'éĂ' 19 | ) 20 | 21 | def test_on_decode_error_show_string(self): 22 | data = MagicMock(spec=six.binary_type) 23 | data.decode.side_effect = UnicodeDecodeError( 24 | 'hitchhiker', 25 | b"", 26 | 42, 27 | 43, 28 | 'the universe and everything else' 29 | ) 30 | with self.assertRaises(UnicodeDecodeError): 31 | self.protocol._decode(data) 32 | 33 | def test_single_packet(self): 34 | self.protocol.add_data( 35 | b'CONNECT\n' 36 | b'accept-version:1.0\n\n\x00' 37 | ) 38 | 39 | frames = self.protocol.pop_frames() 40 | 41 | self.assertEqual(len(frames), 1) 42 | self.assertEqual(frames[0].command, u'CONNECT') 43 | self.assertEqual(frames[0].headers, {u'accept-version': u'1.0'}) 44 | self.assertEqual(frames[0].body, None) 45 | 46 | self.assertEqual(self.protocol._pending_parts, []) 47 | 48 | def test_parcial_packet(self): 49 | stream_data = ( 50 | b'CONNECT\n', 51 | b'accept-version:1.0\n\n\x00', 52 | ) 53 | 54 | for data in stream_data: 55 | self.protocol.add_data(data) 56 | 57 | frames = self.protocol.pop_frames() 58 | 59 | self.assertEqual(len(frames), 1) 60 | self.assertEqual(frames[0].command, u'CONNECT') 61 | self.assertEqual(frames[0].headers, {u'accept-version': u'1.0'}) 62 | self.assertEqual(frames[0].body, None) 63 | 64 | def test_multi_parcial_packet1(self): 65 | stream_data = ( 66 | b'CONNECT\n', 67 | b'accept-version:1.0\n\n\x00\n', 68 | b'CONNECTED\n', 69 | b'version:1.0\n\n\x00\n' 70 | ) 71 | 72 | for data in stream_data: 73 | self.protocol.add_data(data) 74 | 75 | frames = self.protocol.pop_frames() 76 | self.assertEqual(len(frames), 2) 77 | 78 | self.assertEqual(frames[0].command, u'CONNECT') 79 | self.assertEqual(frames[0].headers, {u'accept-version': u'1.0'}) 80 | self.assertEqual(frames[0].body, None) 81 | 82 | self.assertEqual(frames[1].command, u'CONNECTED') 83 | self.assertEqual(frames[1].headers, {u'version': u'1.0'}) 84 | self.assertEqual(frames[1].body, None) 85 | 86 | self.assertEqual(self.protocol._pending_parts, []) 87 | 88 | def test_multi_parcial_packet2(self): 89 | stream_data = ( 90 | b'CONNECTED\n' 91 | b'version:1.0\n\n', 92 | b'\x00\nERROR\n', 93 | b'header:1.0\n\n', 94 | b'Hey dude\x00\n', 95 | ) 96 | 97 | for data in stream_data: 98 | self.protocol.add_data(data) 99 | 100 | frames = self.protocol.pop_frames() 101 | self.assertEqual(len(frames), 2) 102 | 103 | self.assertEqual(frames[0].command, u'CONNECTED') 104 | self.assertEqual(frames[0].headers, {u'version': u'1.0'}) 105 | self.assertEqual(frames[0].body, None) 106 | 107 | self.assertEqual(frames[1].command, u'ERROR') 108 | self.assertEqual(frames[1].headers, {u'header': u'1.0'}) 109 | self.assertEqual(frames[1].body, u'Hey dude') 110 | 111 | self.assertEqual(self.protocol._pending_parts, []) 112 | 113 | def test_multi_parcial_packet_with_utf8(self): 114 | stream_data = ( 115 | b'CONNECTED\n' 116 | b'accept-version:1.0\n\n', 117 | b'\x00\nERROR\n', 118 | b'header:1.0\n\n\xc3', 119 | b'\xa7\x00\n', 120 | ) 121 | 122 | for data in stream_data: 123 | self.protocol.add_data(data) 124 | 125 | self.assertEqual(len(self.protocol._frames_ready), 2) 126 | self.assertEqual(self.protocol._pending_parts, []) 127 | 128 | self.assertEqual(self.protocol._frames_ready[0].body, None) 129 | self.assertEqual(self.protocol._frames_ready[1].body, u'ç') 130 | 131 | def test_heart_beat_packet1(self): 132 | self.protocol._recv_heart_beat = MagicMock() 133 | self.protocol.add_data(b'\n') 134 | 135 | self.assertEqual(self.protocol._pending_parts, []) 136 | self.assertTrue(self.protocol._recv_heart_beat.called) 137 | 138 | def test_heart_beat_packet2(self): 139 | self.protocol._recv_heart_beat = MagicMock() 140 | self.protocol.add_data( 141 | b'CONNECT\n' 142 | b'accept-version:1.0\n\n\x00\n' 143 | ) 144 | 145 | self.assertTrue(self.protocol._recv_heart_beat.called) 146 | self.assertEqual(self.protocol._pending_parts, []) 147 | 148 | def test_heart_beat_packet3(self): 149 | self.protocol._recv_heart_beat = MagicMock() 150 | self.protocol.add_data( 151 | b'\nCONNECT\n' 152 | b'accept-version:1.0\n\n\x00' 153 | ) 154 | 155 | frames = self.protocol.pop_frames() 156 | self.assertEqual(len(frames), 1) 157 | 158 | self.assertEqual(frames[0].command, u'CONNECT') 159 | self.assertEqual(frames[0].headers, {u'accept-version': u'1.0'}) 160 | self.assertEqual(frames[0].body, None) 161 | 162 | self.assertTrue(self.protocol._recv_heart_beat.called) 163 | self.assertEqual(self.protocol._pending_parts, []) 164 | 165 | 166 | class TestBuildFrame(TestCase): 167 | 168 | def setUp(self): 169 | self.protocol = StompProtocol() 170 | 171 | def test_build_frame_with_body(self): 172 | buf = self.protocol.build_frame('HELLO', { 173 | 'from': 'me', 174 | 'to': 'you' 175 | }, 'I Am The Walrus') 176 | 177 | self.assertEqual( 178 | buf, 179 | b'HELLO\n' 180 | b'from:me\n' 181 | b'to:you\n\n' 182 | b'I Am The Walrus' 183 | b'\x00') 184 | 185 | def test_build_frame_without_body(self): 186 | buf = self.protocol.build_frame('HI', { 187 | 'from': '1', 188 | 'to': '2' 189 | }) 190 | 191 | self.assertEqual( 192 | buf, 193 | b'HI\n' 194 | b'from:1\n' 195 | b'to:2\n\n' 196 | b'\x00') 197 | 198 | 199 | class TestReadFrame(TestCase): 200 | 201 | def setUp(self): 202 | self.protocol = StompProtocol() 203 | 204 | def test_single_packet(self): 205 | self.protocol.add_data( 206 | b'CONNECT\n' 207 | b'accept-version:1.0\n\n\x00' 208 | ) 209 | 210 | self.assertEqual(len(self.protocol._frames_ready), 1) 211 | 212 | frame = self.protocol._frames_ready[0] 213 | self.assertEqual(frame.command, 'CONNECT') 214 | self.assertEqual(frame.headers, {'accept-version': '1.0'}) 215 | self.assertEqual(frame.body, None) 216 | -------------------------------------------------------------------------------- /torstomp/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | import socket 3 | import logging 4 | import datetime 5 | 6 | from tornado.iostream import IOStream, StreamClosedError 7 | from tornado.ioloop import IOLoop 8 | from tornado import gen 9 | 10 | from datetime import timedelta 11 | 12 | from torstomp.protocol import StompProtocol 13 | from torstomp.errors import StompError 14 | from torstomp.subscription import Subscription 15 | 16 | 17 | class TorStomp(object): 18 | 19 | VERSION = '1.1' 20 | 21 | def __init__(self, host='localhost', port=61613, connect_headers={}, 22 | on_error=None, on_disconnect=None, on_connect=None, 23 | reconnect_max_attempts=-1, reconnect_timeout=1000, 24 | log_name='TorStomp'): 25 | 26 | self.host = host 27 | self.port = port 28 | self.logger = logging.getLogger(log_name) 29 | 30 | self._connect_headers = connect_headers 31 | self._connect_headers['accept-version'] = self.VERSION 32 | self._heart_beat_handler = None 33 | self.connected = False 34 | self.disconnected_date = None 35 | self._disconnecting = False 36 | self._protocol = StompProtocol(log_name=log_name) 37 | self._subscriptions = {} 38 | self._last_subscribe_id = 0 39 | self._on_error = on_error 40 | self._on_disconnect = on_disconnect 41 | self._on_connect = on_connect 42 | 43 | self._reconnect_max_attempts = reconnect_max_attempts 44 | self._reconnect_timeout = timedelta(milliseconds=reconnect_timeout) 45 | self._reconnect_attempts = 0 46 | 47 | @gen.coroutine 48 | def connect(self): 49 | self.stream = self._build_io_stream() 50 | 51 | try: 52 | yield self.stream.connect((self.host, self.port)) 53 | self.logger.info('Stomp connection estabilished') 54 | except socket.error as error: 55 | self.logger.error( 56 | '[attempt: %d] Connect error: %s', self._reconnect_attempts, 57 | error) 58 | self._schedule_reconnect() 59 | return 60 | 61 | self.stream.set_close_callback(self._on_disconnect_socket) 62 | self.stream.read_until_close( 63 | streaming_callback=self._on_data, 64 | callback=self._on_data) 65 | 66 | self.connected = True 67 | self._disconnecting = False 68 | self._reconnect_attempts = 0 69 | self._protocol.reset() 70 | 71 | yield self._send_frame('CONNECT', self._connect_headers) 72 | 73 | for subscription in self._subscriptions.values(): 74 | yield self._send_subscribe_frame(subscription) 75 | 76 | if self._on_connect: 77 | self._on_connect() 78 | 79 | def subscribe(self, destination, ack='auto', extra_headers={}, 80 | callback=None): 81 | 82 | self._last_subscribe_id += 1 83 | 84 | subscription = Subscription( 85 | destination=destination, 86 | id=self._last_subscribe_id, 87 | ack=ack, 88 | extra_headers=extra_headers, 89 | callback=callback) 90 | 91 | self._subscriptions[str(self._last_subscribe_id)] = subscription 92 | 93 | if self.connected: 94 | self._send_subscribe_frame(subscription) 95 | 96 | def unsubscribe(self, subscription): 97 | subscription_id = str(subscription.id) 98 | 99 | if subscription_id in self._subscriptions.keys(): 100 | self._send_unsubscribe_frame(subscription) 101 | del self._subscriptions[subscription_id] 102 | 103 | def send(self, destination, body='', headers={}, send_content_length=True): 104 | headers['destination'] = destination 105 | 106 | if body: 107 | body = self._protocol._encode(body) 108 | 109 | # ActiveMQ determines the type of a message by the 110 | # inclusion of the content-length header 111 | if send_content_length: 112 | headers['content-length'] = len(body) 113 | 114 | return self._send_frame('SEND', headers, body) 115 | 116 | def ack(self, frame): 117 | headers = { 118 | 'subscription': frame.headers['subscription'], 119 | 'message-id': frame.headers['message-id'] 120 | } 121 | 122 | return self._send_frame('ACK', headers) 123 | 124 | def nack(self, frame): 125 | headers = { 126 | 'subscription': frame.headers['subscription'], 127 | 'message-id': frame.headers['message-id'] 128 | } 129 | 130 | return self._send_frame('NACK', headers) 131 | 132 | def _build_io_stream(self): 133 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) 134 | return IOStream(s) 135 | 136 | def _on_disconnect_socket(self): 137 | self._stop_scheduled_heart_beat() 138 | self.connected = False 139 | self.disconnected_date = datetime.datetime.now() 140 | 141 | if self._disconnecting: 142 | self.logger.info('TCP connection end gracefully') 143 | else: 144 | self.logger.info('TCP connection unexpected end') 145 | self._schedule_reconnect() 146 | 147 | if self._on_disconnect: 148 | self._on_disconnect() 149 | 150 | def _schedule_reconnect(self): 151 | if self._reconnect_max_attempts == -1 or \ 152 | self._reconnect_attempts < self._reconnect_max_attempts: 153 | 154 | self._reconnect_attempts += 1 155 | self._reconnect_timeout_handler = IOLoop.current().add_timeout( 156 | self._reconnect_timeout, self.connect) 157 | else: 158 | self.logger.error('All Connection attempts failed') 159 | 160 | def _on_data(self, data): 161 | if not data: 162 | return 163 | 164 | self._protocol.add_data(data) 165 | 166 | frames = self._protocol.pop_frames() 167 | if frames: 168 | self._received_frames(frames) 169 | 170 | def _send_frame(self, command, headers={}, body=''): 171 | buf = self._protocol.build_frame(command, headers, body) 172 | return self.stream.write(buf) 173 | 174 | def _set_connected(self, connected_frame): 175 | heartbeat = connected_frame.headers.get('heart-beat') 176 | 177 | if heartbeat: 178 | sx, sy = heartbeat.split(',') 179 | sx, sy = int(sx), int(sy) 180 | 181 | if sy: 182 | self._set_heart_beat(sy) 183 | 184 | def _set_heart_beat(self, time): 185 | self._heart_beat_delta = timedelta(milliseconds=time) 186 | self._stop_scheduled_heart_beat() 187 | 188 | self._do_heart_beat() 189 | 190 | def _schedule_heart_beat(self): 191 | self._heart_beat_handler = IOLoop.current().add_timeout( 192 | self._heart_beat_delta, self._do_heart_beat) 193 | 194 | def _stop_scheduled_heart_beat(self): 195 | if self._heart_beat_handler: 196 | IOLoop.current().remove_timeout(self._heart_beat_handler) 197 | 198 | self._heart_beat_handler = None 199 | 200 | def _do_heart_beat(self): 201 | self.logger.debug('Sending heartbeat') 202 | 203 | try: 204 | self.stream.write(self._protocol.HEART_BEAT) 205 | except StreamClosedError: 206 | logging.warning('Heart beat failed: stream is closed') 207 | 208 | self._schedule_heart_beat() 209 | 210 | def _received_frames(self, frames): 211 | for frame in frames: 212 | if frame.command == 'MESSAGE': 213 | self._received_message_frame(frame) 214 | elif frame.command == 'CONNECTED': 215 | self._set_connected(frame) 216 | elif frame.command == 'ERROR': 217 | self._received_error_frame(frame) 218 | else: 219 | self._received_unhandled_frame(frame) 220 | 221 | def _received_message_frame(self, frame): 222 | subscription_header = frame.headers.get('subscription') 223 | 224 | subscription = self._subscriptions.get(subscription_header) 225 | 226 | if not subscription: 227 | self.logger.error( 228 | 'Not found subscription %d' % subscription_header) 229 | return 230 | 231 | subscription.callback(frame, frame.body) 232 | 233 | def _received_error_frame(self, frame): 234 | message = frame.headers.get('message') 235 | 236 | self.logger.error('Received error: %s', message) 237 | self.logger.debug('Error detail %s', frame.body) 238 | 239 | if self._on_error: 240 | self._on_error( 241 | StompError(message, frame.body)) 242 | 243 | def _received_unhandled_frame(self, frame): 244 | self.logger.warn('Received unhandled frame: %s', frame.command) 245 | 246 | def _send_subscribe_frame(self, subscription): 247 | headers = { 248 | 'id': subscription.id, 249 | 'destination': subscription.destination, 250 | 'ack': subscription.ack 251 | } 252 | headers.update(subscription.extra_headers) 253 | 254 | return self._send_frame('SUBSCRIBE', headers) 255 | 256 | def _send_unsubscribe_frame(self, subscription): 257 | headers = { 258 | 'id': subscription.id, 259 | 'destination': subscription.destination 260 | } 261 | return self._send_frame('UNSUBSCRIBE', headers) 262 | -------------------------------------------------------------------------------- /torstomp/errors.py: -------------------------------------------------------------------------------- 1 | 2 | class StompError(Exception): 3 | 4 | def __init__(self, message, detail): 5 | super(StompError, self).__init__(message) 6 | self.detail = detail 7 | -------------------------------------------------------------------------------- /torstomp/frame.py: -------------------------------------------------------------------------------- 1 | class Frame(object): 2 | 3 | def __init__(self, command, headers, body): 4 | self.command = command 5 | self.headers = headers 6 | self.body = body 7 | 8 | def __repr__(self): 9 | return '' % self.command 10 | -------------------------------------------------------------------------------- /torstomp/protocol.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | import logging 3 | import sys 4 | import six 5 | 6 | from torstomp.frame import Frame 7 | 8 | PYTHON3 = sys.hexversion >= 0x03000000 9 | 10 | if not PYTHON3: 11 | import codecs 12 | utf8_decoder = codecs.lookup('utf-8') 13 | 14 | 15 | class StompProtocol(object): 16 | 17 | HEART_BEAT = b'\n' 18 | EOF = b'\x00' 19 | 20 | def __init__(self, log_name='StompProtocol'): 21 | self._pending_parts = [] 22 | self._frames_ready = [] 23 | self.logger = logging.getLogger(log_name) 24 | 25 | def _decode(self, byte_data): 26 | try: 27 | if isinstance(byte_data, six.binary_type): 28 | return byte_data.decode('utf-8') 29 | 30 | return byte_data 31 | except UnicodeDecodeError: 32 | logging.error(u"string was: {}".format(byte_data)) 33 | raise 34 | 35 | def _encode(self, value): 36 | if isinstance(value, six.text_type): 37 | return value.encode('utf-8') 38 | 39 | return value 40 | 41 | def reset(self): 42 | self._pending_parts = [] 43 | self._frames_ready = [] 44 | 45 | def add_data(self, data): 46 | if not self._pending_parts and data.startswith(self.HEART_BEAT): 47 | self._recv_heart_beat() 48 | data = data[1:] 49 | 50 | if data: 51 | return self.add_data(data) 52 | 53 | before_eof, sep, after_eof = data.partition(self.EOF) 54 | 55 | if before_eof: 56 | self._pending_parts.append(before_eof) 57 | 58 | if sep: 59 | frame_data = b''.join(self._pending_parts) 60 | self._pending_parts = [] 61 | self._proccess_frame(frame_data) 62 | 63 | if after_eof: 64 | self.add_data(after_eof) 65 | 66 | def _proccess_frame(self, data): 67 | data = self._decode(data) 68 | command, remaing = data.split('\n', 1) 69 | 70 | raw_headers, remaing = remaing.split('\n\n', 1) 71 | headers = dict([l.split(':', 1) for l in raw_headers.split('\n')]) 72 | body = remaing if remaing else None 73 | 74 | self._frames_ready.append(Frame(command, headers=headers, body=body)) 75 | 76 | def _recv_heart_beat(self): 77 | self.logger.debug('Heartbeat received') 78 | 79 | def build_frame(self, command, headers={}, body=''): 80 | lines = [command, '\n'] 81 | 82 | for key, value in sorted(headers.items()): 83 | lines.append('%s:%s\n' % (key, value)) 84 | 85 | lines.append('\n') 86 | lines.append(body) 87 | lines.append(self.EOF) 88 | 89 | return b''.join([self._encode(line) for line in lines]) 90 | 91 | def pop_frames(self): 92 | frames = self._frames_ready 93 | self._frames_ready = [] 94 | 95 | return frames 96 | -------------------------------------------------------------------------------- /torstomp/subscription.py: -------------------------------------------------------------------------------- 1 | class Subscription(object): 2 | 3 | def __init__(self, destination, id, ack, extra_headers, callback): 4 | self.destination = destination 5 | self.id = id 6 | self.ack = ack 7 | self.extra_headers = extra_headers 8 | self.callback = callback 9 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [flake8] 2 | exclude = ./build/* 3 | max-line-length = 130 4 | max-complexity = 20 5 | --------------------------------------------------------------------------------