├── .gitignore ├── .travis.yml ├── CHANGELOG ├── DOCUMENTATION.rst ├── LICENSE.txt ├── MANIFEST.in ├── README.rst ├── development.txt ├── docs ├── amqp-xml-doc0-9-1.pdf └── amqp0-9-1.pdf ├── examples ├── gevent_hello.py ├── rpc_client.py ├── rpc_server.py ├── synchronous └── synchronous_gevent ├── haigha ├── __init__.py ├── channel.py ├── channel_pool.py ├── classes │ ├── __init__.py │ ├── basic_class.py │ ├── channel_class.py │ ├── exchange_class.py │ ├── protocol_class.py │ ├── queue_class.py │ └── transaction_class.py ├── connection.py ├── connections │ ├── __init__.py │ └── rabbit_connection.py ├── exceptions.py ├── frames │ ├── __init__.py │ ├── content_frame.py │ ├── frame.py │ ├── header_frame.py │ ├── heartbeat_frame.py │ └── method_frame.py ├── message.py ├── reader.py ├── transports │ ├── __init__.py │ ├── event_transport.py │ ├── gevent_transport.py │ ├── socket_transport.py │ └── transport.py └── writer.py ├── requirements.txt ├── scripts ├── console ├── rabbit_table_test ├── regression │ ├── issue_24 │ ├── issue_31 │ └── issue_33 ├── stress_test └── test ├── setup.py └── tests ├── __init__.py ├── integration ├── __init__.py ├── channel_basic_test.py └── rabbit_extensions_test.py └── unit ├── __init__.py ├── channel_pool_test.py ├── channel_test.py ├── classes ├── __init__.py ├── basic_class_test.py ├── channel_class_test.py ├── exchange_class_test.py ├── protocol_class_test.py ├── queue_class_test.py └── transaction_class_test.py ├── connection_test.py ├── connections └── rabbit_connection_test.py ├── frames ├── __init__.py ├── content_frame_test.py ├── frame_test.py ├── header_frame_test.py ├── heartbeat_frame_test.py └── method_frame_test.py ├── message_test.py ├── reader_test.py ├── transports ├── __init__.py ├── event_transport_test.py ├── gevent_transport_test.py ├── socket_transport_test.py └── transport_test.py └── writer_test.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.swp 3 | *.egg-info 4 | .coverage 5 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | #- "3.4" 4 | #- "3.3" 5 | #- "3.2" 6 | - "2.7" 7 | - "pypy" 8 | script: nosetests 9 | install: 10 | - "pip install -r development.txt" 11 | - "if [ ${TRAVIS_PYTHON_VERSION} != 'pypy' ]; then pip install gevent; fi" 12 | - "pip install nose" 13 | services: 14 | - rabbitmq 15 | -------------------------------------------------------------------------------- /CHANGELOG: -------------------------------------------------------------------------------- 1 | 0.9.0 2 | ===== 3 | 4 | Merge https://github.com/agoragames/haigha/pull/69 from @kevinconway 5 | Add support for IPv6 connections. Cycle through all available addresses 6 | from socket.getaddrinfo() in their natural order. 7 | 8 | Merge https://github.com/agoragames/haigha/pull/74 from @vitaly-krugl 9 | Significant cleanup of channel frame processing and error handling. 10 | 11 | Merge https://github.com/agoragames/haigha/pull/82 from @vitaly-krugl 12 | Implemented support for RabbitMQ broker-initiated Basic.Cancel. 13 | 14 | Integration test framework. Integration tests running on Travis CI. 15 | 16 | 17 | 0.8.0 18 | ===== 19 | 20 | Merge https://github.com/agoragames/haigha/pull/73 from @vitaly-krugl 21 | Adds listener support for basic.return, fixes protocol errors when not 22 | correctly reading the list of frames upon a basic.return. 23 | 24 | Merge https://github.com/agoragames/haigha/pull/78 from @liorn 25 | Correctly handle heartbeat intervals and disconnect when no traffic according 26 | to 0.9.1 spec section 4.2.7 27 | 28 | 0.7.3 29 | ===== 30 | 31 | Merge https://github.com/agoragames/haigha/pull/66 from @rocco66 32 | Adds IPv6 support. 33 | 34 | 0.7.2 35 | ===== 36 | 37 | Merge https://github.com/agoragames/haigha/pull/64 from @subzey 38 | Fixes installation on Windows. 39 | 40 | 0.7.1 41 | ===== 42 | 43 | PEP8 styling fixes 44 | 45 | Merge https://github.com/agoragames/haigha/pull/55 from @dizpers 46 | Adds keyword argument passing to GeventPoolTransport constructor 47 | from the Connection constructor. 48 | 49 | 0.7.0 50 | ===== 51 | 52 | Adds `synchronous` keyword argument to `Connection` constructor, which makes 53 | the connection behave as if it were synchronous, no matter the underlying 54 | transport type. Useful for gevent where one does not want to use callbacks. 55 | Applies to all channels, enforces `nowait=False` where applicable. 56 | 57 | Adds `synchronous` keyword argument to `Connection.channel`, which makes the 58 | `Channel` object synchronous independent of the underlying transport type. 59 | For example, if a synchronous channel on an asynchronous transport has a 60 | protocol method called with `nowait=False` (where applicable), then the method 61 | will not return until the response frame has been read. Does not enforce 62 | `nowait=False` where applicable. 63 | 64 | Adds `synchronous_connect` option to Connection constructor which will enforce 65 | synchronous behavior in `Connection.connect` regardless of the underlying 66 | transport type. Improves features to handle issue #7. Also makes 67 | `Connection.close` synchronous. Defaults to True if transport is synchronous 68 | or `synchronous=True` passed in constructor. 69 | https://github.com/agoragames/haigha/issues/7 70 | 71 | Standard and gevent socket transport will raise EnvironmentErrors that aren't 72 | in `(errno.EAGAIN,errno.EWOULDBLOCK,errno.EINTR)`. Fixes notifying read loops 73 | of potential problems and fixes #44 74 | https://github.com/agoragames/haigha/issues/44 75 | 76 | Immediately close connection and raise ConnectionClosed if there is a 77 | FrameError in reading frames from the broker in `Connection.read_frames`. 78 | 79 | Detect frames that are larger than negotiated frame max in 80 | `Connection.send_frame`, immediately close connection and raise 81 | ConnectionClosed. 82 | 83 | Add property `Connection.closed` 84 | 85 | 0.6.2 86 | ===== 87 | 88 | Raise ConnectionClosed in Connection.synchronous if the transport has 89 | been cleared 90 | 91 | 0.6.1 92 | ===== 93 | 94 | Merge https://github.com/agoragames/haigha/pull/35 from @xjdrew 95 | Fix exception raising when there's an application error on a synchronous transport 96 | Fixes synchronous return values from queue.declare 97 | Adds an example synchronous RPC client and server 98 | 99 | 0.6.0 100 | ===== 101 | 102 | Fixes #33, only accept (str,unicode,bytearray) as Message body types. Default 103 | to empty string. 104 | 105 | 0.5.12 106 | ====== 107 | 108 | Fixes #31 wherein a channel is closed on a synchronous transport while reading 109 | frames and waiting for the synchronous callback to be executed. 110 | 111 | 0.5.11 112 | ====== 113 | 114 | Fixed writing headers if the type of a field is a subclass of the supported 115 | native types 116 | 117 | 0.5.10 118 | ====== 119 | 120 | Fixed unexpected indent error. 121 | 122 | 0.5.9 123 | ===== 124 | 125 | Support array lists in headers. Thanks to @joeyimbasciano 126 | https://github.com/agoragames/haigha/pull/28 127 | 128 | 129 | 0.5.8 130 | ===== 131 | 132 | Merge pull requests, suppresses DeprecationWarning in latest gevent. 133 | 134 | 0.5.7 135 | ===== 136 | 137 | Fixed exception in handling in SocketTransport when sendall() raises a socket 138 | error. Also fixed re-raise of exception on SocketTransport.read() so that it 139 | always marks the transport as closed. 140 | 141 | 0.5.6 142 | ===== 143 | 144 | Only catch Exception subclasses in Connection and Channel. Fixes issue #23. 145 | 146 | Update to Chai 0.2.0 147 | 148 | 149 | 0.5.5 150 | ===== 151 | 152 | Fixed bug in gevent transport connect parameter to socket transport 153 | 154 | 155 | 0.5.4 156 | ===== 157 | 158 | Gevent transport can be used without monkey patching sockets 159 | 160 | If there is an error in a channel close callback, don't fail in the second 161 | call to Channel.close 162 | 163 | 164 | 0.5.3 165 | ===== 166 | 167 | Added 'haigha.connections' to package. Thanks to @grncdr 168 | https://github.com/agoragames/haigha/pull/16 169 | 170 | 171 | 0.5.2 172 | ===== 173 | 174 | Added haigha.connections.RabbitConnection to support RabbitMQ-specific 175 | features. This includes `auto_delete` and `internal` parameters to 176 | exchange.declare, exchange.bind, exchange.unbind, client side basic.nack, 177 | and confirm.select to enable publisher confirms. See basic.set_ack_listener 178 | and basic.set_nack_listener for handling the confirmations. 179 | http://www.rabbitmq.com/blog/2011/02/10/introducing-publisher-confirms/ 180 | 181 | Removed `auto_delete` and `internal` parameters from common exchange.declare 182 | 183 | Default transport is now the SocketTransport. This is unlikely to change again. 184 | 185 | Implemented PEP 396. 186 | 187 | 188 | 0.5.1 189 | ===== 190 | 191 | GeventTransport subclasses SocketTransport 192 | 193 | In all cases, if a method has 'nowait' argument and an optional callback is 194 | supplied, the method will execute as 'nowait=False', i.e. synchronously 195 | 196 | Connection supplies Channels with the protocol mapping 197 | 198 | Connection close_info is available if the connection is closed or disconnected 199 | 200 | The reply_text in Channel.close is truncated to 255 characters when sent to 201 | the broker 202 | 203 | Fixed bug in Channel.clear_synchronous_cb which needs to return the argument 204 | if no matching callback was found. This fixes a bug in ProtocolClass.dispatch. 205 | 206 | Fix isinstance parameter ordering in Channel.SyncWrapper.__eq__ 207 | 208 | 209 | 0.5.0 210 | ===== 211 | 212 | Fix message reading in basic.get 213 | 214 | Added optional open_cb kwarg to Connection constructor 215 | 216 | Added optional callback to basic.consume for notifications when broker has 217 | registered the consumer 218 | 219 | Moved channel state out of ChannelClass and into Channel to fix access problems 220 | after Channel has been cleaned up 221 | 222 | Added support for Channel open notification listeners 223 | 224 | All AMQP timestamps are in UTC 225 | 226 | Most exceptions will now propagate to user code, fixing problems with 227 | gevent.GreenletExit and SystemExit 228 | 229 | Preliminary support for synchronous clients 230 | 231 | 232 | 0.4.5 233 | ===== 234 | 235 | Fix heartbeat support in event transport 236 | 237 | 238 | 0.4.4 239 | ===== 240 | 241 | Fix handling of socket timeouts in gevent transport 242 | 243 | 244 | 0.4.3 245 | ===== 246 | 247 | Client sends heartbeat frames on its own rather than simply replying to ones 248 | received from the broker 249 | 250 | 251 | 0.4.2 252 | ===== 253 | 254 | Channel close listeners are now notified before channel local data is cleaned 255 | up 256 | 257 | Default to the gevent transport 258 | 259 | Removed fixed requirements for libevent or gevent. It is now up to the user to 260 | install the packages corresponding to their preferred transport. 261 | 262 | 263 | 0.4.1 264 | ===== 265 | 266 | Added a write lock around GeventTransport.write() call to socket.sendall() 267 | 268 | 269 | 0.4.0 270 | ===== 271 | 272 | Fixed handling of channel closures due to error. This cleaned up several DoS 273 | vulnerabilities, including a lockup in the ChannelPool, unbounded memory 274 | consumption, and failure to re-use channel ids after max id enumerator 275 | rolled over. Channel cleanup relies on reference counting to ensure a client 276 | can respond to rapid opening and closing of channels. 277 | 278 | Fixed failure to always cast GeventTransport.read() return to a bytearray, 279 | which caused problems when reading partial frames. 280 | 281 | Fixed failure in GeventTransport to switch to non-blocking socket after 282 | connecting. 283 | 284 | Fixed fetching of close_info when connection close callback executes. 285 | 286 | 287 | 0.3.4 288 | ===== 289 | 290 | Abstracted libevent requirement and added support for gevent backend 291 | 292 | Support non-blocking connect on libevent transport 293 | 294 | Removed connection strategies from Connection and deleted ConnectionStrategy 295 | class. Combined with a simplified and abstracted approach to IO and scheduling 296 | in the Transports, the strategy for failed connections is more deterministic 297 | and left in the hands of the user. 298 | 299 | 300 | 0.3.3 301 | ===== 302 | 303 | Added cap to ChannelPool. This should be used to ensure that all of the broker 304 | memory is consumed by channels in cases where it is slow to process 305 | transactions. 306 | 307 | 0.3.2 308 | ===== 309 | 310 | Fixed calling into user-provided close callback, changed interface of close_cb 311 | to accept no arguments 312 | 313 | Use buffer-per-frame in Connection.send_frame() to work around a memory 314 | freeing problem in EventSocket 315 | 316 | First round of in-depth documentation in DOCUMENTATION.rst 317 | 318 | Added callbacks to Exchange.declare(), Exchange.delete() and 319 | Transaction.select(). 320 | 321 | Fixed FIFO ordering problems for all methods that support callbacks to user 322 | after a synchronous operation. 323 | 324 | 325 | 0.3.1 326 | ===== 327 | 328 | Removed encoding of arrays as table fields 329 | 330 | 331 | 0.3.0 332 | ===== 333 | 334 | Test coverage of all protocol and frame classes 335 | 336 | Test coverage, bug fixes and minor performance improvements to Reader and 337 | Writer classes 338 | 339 | Reader and Writer conform to RabbitMQ and QPid errata as documented 340 | at http://dev.rabbitmq.com/wiki/Amqp091Errata#section_3. Added 341 | scripts/rabbit_table_test for integration testing. 342 | 343 | 344 | 0.2.3 345 | ===== 346 | 347 | First changelog. 348 | 349 | Removed manipulation of channels in Connection.disconnect(). It's now 350 | up to the application to manage channels when explicitly shutting down a 351 | connection or when the socket is dropped and the strategy reconnects. This 352 | fixed a problem wherein an application may have created a connection, 353 | created some channels and queued operations on them, but it takes the 354 | strategy a few attempts before completing a socket connection. 355 | 356 | Fixed character escpaing when logging outbound ContentFrames. 357 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 2 | 3 | Redistribution and use in source and binary forms, with or without 4 | modification, are permitted provided that the following conditions 5 | are met: 6 | 7 | * Redistributions of source code must retain the above copyright notice, 8 | this list of conditions and the following disclaimer. 9 | 10 | * Redistributions in binary form must reproduce the above copyright 11 | notice, this list of conditions and the following disclaimer in the 12 | documentation and/or other materials provided with the distribution. 13 | 14 | * Neither the name of Agora Games nor the names of its contributors 15 | may be used to endorse or promote products derived from this software 16 | without specific prior written permission. 17 | 18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 19 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 20 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 21 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 22 | HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 23 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED 24 | TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR 25 | PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF 26 | LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING 27 | NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 28 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include *.txt 2 | include README.rst 3 | include MANIFEST.in 4 | recursive-include scripts * 5 | recursive-include haigha *.py 6 | prune *.pyc 7 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | ========================================================== 2 | Haigha - Synchronous and asynchronous AMQP client library 3 | ========================================================== 4 | 5 | .. image:: https://travis-ci.org/agoragames/haigha.svg?branch=master 6 | :target: https://travis-ci.org/agoragames/haigha 7 | 8 | 9 | :Version: 0.9.0 10 | :Download: http://pypi.python.org/pypi/haigha 11 | :Source: https://github.com/agoragames/haigha 12 | :Keywords: python, amqp, rabbitmq, event, libevent, gevent 13 | 14 | .. contents:: 15 | :local: 16 | 17 | .. _haigha-overview: 18 | 19 | Overview 20 | ======== 21 | 22 | Haigha provides a simple to use client library for interacting with AMQP brokers. It currently supports the 0.9.1 protocol and is integration tested against the latest RabbitMQ 2.8.1 (see `errata `_). Haigha is a descendant of ``py-amqplib`` and owes much to its developers. 23 | 24 | The goals of haigha are performance, simplicity, and adherence to the form and function of the AMQP protocol. It adds a few useful features, such as the ``ChannelPool`` class and ``Channel.publish_synchronous``, to ease use of powerful features in real-world applications. 25 | 26 | By default, Haigha operates in a completely asynchronous mode, relying on callbacks to notify application code of responses from the broker. Where applicable, ``nowait`` defaults to ``True``. The application code is welcome to call a series of methods, and Haigha will manage the stack and synchronous handshakes in the event loop. 27 | 28 | Starting with the 0.5.0 series, haigha natively supports 3 transport types; libevent, gevent and standard sockets. The socket implementation defaults to synchronous mode and is useful for an interactive console or scripting, and the gevent transport is the preferred asynchronous backend though it can also be used synchronously as well. 29 | 30 | Documentation 31 | ============= 32 | 33 | This file and the various files in the ``scripts`` directory serve as a simple introduction to haigha. For more complete documentation, see `DOCUMENTATION.rst `_. 34 | 35 | 36 | Example 37 | ======= 38 | 39 | See the ``scripts`` and ``examples`` directories for several examples, in particular the ``stress_test`` script which you can use to test the performance of haigha against your broker. Below is a simple example of a client that connects, processes one message and quits. :: 40 | 41 | from haigha.connection import Connection 42 | from haigha.message import Message 43 | 44 | connection = Connection( 45 | user='guest', password='guest', 46 | vhost='/', host='localhost', 47 | heartbeat=None, debug=True) 48 | 49 | ch = connection.channel() 50 | ch.exchange.declare('test_exchange', 'direct') 51 | ch.queue.declare('test_queue', auto_delete=True) 52 | ch.queue.bind('test_queue', 'test_exchange', 'test_key') 53 | ch.basic.publish( Message('body', application_headers={'hello':'world'}), 54 | 'test_exchange', 'test_key' ) 55 | print ch.basic.get('test_queue') 56 | connection.close() 57 | 58 | To use protocol extensions for RabbitMQ, initialize the connection with the ``haigha.connections.rabbit_connection.RabbitConnection`` class. 59 | 60 | Roadmap 61 | ======= 62 | 63 | * Documentation (there's always more) 64 | * Improved error handling 65 | * Implementation of error codes in the spec 66 | * Testing and integration with brokers other than RabbitMQ 67 | * Identify and improve inefficient code 68 | * Edge cases in frame management 69 | * Improvements to usabililty 70 | * SSL 71 | * Allow nowait when asynchronous transport but Connection put into synchronous mode. 72 | 73 | Haigha has been tested exclusively with Python 2.6 and 2.7, but we intend for it to work with the 3.x series as well. Please `report `_ any issues you may have. 74 | 75 | Installation 76 | ============ 77 | 78 | Haigha is available on `pypi `_ and can be installed using ``pip`` :: 79 | 80 | pip install haigha 81 | 82 | If installing from source: 83 | 84 | * with development requirements (e.g. testing frameworks) :: 85 | 86 | pip install -r development.txt 87 | 88 | * without development requirements :: 89 | 90 | pip install -r requirements.txt 91 | 92 | Note that haigha does not install either gevent or libevent support automatically. For libevent, haigha has been tested and deployed with the ``event-agora==0.4.1`` library. 93 | 94 | 95 | Testing 96 | ======= 97 | 98 | Unit tests can be run with either the included script, or with `nose `_ :: 99 | 100 | ./haigha$ scripts/test 101 | ./haigha$ nosetests 102 | 103 | There are two other testing scripts of note. ``rabbit_table_test`` is a simple integration test that confirms compliance with RabbitMQ `errata `_. The ``stress_test`` script is a valuable tool that offers load-testing capability similar to `Apache Bench `_ or `Siege `_. It is used both to confirm the robustness of haigha, as well as benchmark hardware or a broker configuration. 104 | 105 | Bug tracker 106 | =========== 107 | 108 | If you have any suggestions, bug reports or annoyances please report them 109 | to our issue tracker at https://github.com/agoragames/haigha/issues 110 | 111 | License 112 | ======= 113 | 114 | This software is licensed under the `New BSD License`. See the ``LICENSE.txt`` 115 | file in the top distribution directory for the full license text. 116 | 117 | .. # vim: syntax=rst expandtab tabstop=4 shiftwidth=4 shiftround 118 | -------------------------------------------------------------------------------- /development.txt: -------------------------------------------------------------------------------- 1 | -r requirements.txt 2 | 3 | chai>=0.4.7 4 | unittest2 5 | -------------------------------------------------------------------------------- /docs/amqp-xml-doc0-9-1.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/docs/amqp-xml-doc0-9-1.pdf -------------------------------------------------------------------------------- /docs/amqp0-9-1.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/docs/amqp0-9-1.pdf -------------------------------------------------------------------------------- /examples/gevent_hello.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | """ Demonstrates publishing and receiving a message via Haigha library using 4 | gevent-based transport. 5 | 6 | Assumes AMQP broker (e.g., RabbitMQ) is running on same machine (localhost) 7 | and is configured with default parameters: 8 | user: guest 9 | password: guest 10 | port: 5672 11 | vhost: '/' 12 | """ 13 | import sys, os 14 | sys.path.append(os.path.abspath(".")) 15 | sys.path.append(os.path.abspath("..")) 16 | 17 | import logging 18 | 19 | import gevent 20 | import gevent.event as gevent_event 21 | 22 | from haigha.connection import Connection as haigha_Connection 23 | from haigha.message import Message 24 | 25 | 26 | class HaighaGeventHello(object): 27 | 28 | def __init__(self, done_cb): 29 | self._done_cb = done_cb 30 | 31 | # Connect to AMQP broker with default connection and authentication 32 | # settings (assumes broker is on localhost) 33 | self._conn = haigha_Connection(transport='gevent', 34 | close_cb=self._connection_closed_cb, 35 | logger=logging.getLogger()) 36 | 37 | # Start message pump 38 | self._message_pump_greenlet = gevent.spawn(self._message_pump_greenthread) 39 | 40 | # Create message channel 41 | self._channel = self._conn.channel() 42 | self._channel.add_close_listener(self._channel_closed_cb) 43 | 44 | # Create and configure message exchange and queue 45 | self._channel.exchange.declare('test_exchange', 'direct') 46 | self._channel.queue.declare('test_queue', auto_delete=True) 47 | self._channel.queue.bind('test_queue', 'test_exchange', 'test_routing_key') 48 | self._channel.basic.consume(queue='test_queue', 49 | consumer=self._handle_incoming_messages) 50 | 51 | # Publish a message on the channel 52 | msg = Message('body', application_headers={'hello':'world'}) 53 | print "Publising message: %s" % (msg,) 54 | self._channel.basic.publish(msg, 'test_exchange', 'test_routing_key') 55 | return 56 | 57 | 58 | def _message_pump_greenthread(self): 59 | print "Entering Message Pump" 60 | try: 61 | while self._conn is not None: 62 | # Pump 63 | self._conn.read_frames() 64 | 65 | # Yield to other greenlets so they don't starve 66 | gevent.sleep() 67 | finally: 68 | print "Leaving Message Pump" 69 | self._done_cb() 70 | return 71 | 72 | 73 | def _handle_incoming_messages(self, msg): 74 | print 75 | print "Received message: %s" % (msg,) 76 | print 77 | 78 | # Initiate graceful closing of the channel 79 | self._channel.basic.cancel(consumer=self._handle_incoming_messages) 80 | self._channel.close() 81 | return 82 | 83 | 84 | def _channel_closed_cb(self, ch): 85 | print "AMQP channel closed; close-info: %s" % ( 86 | self._channel.close_info,) 87 | self._channel = None 88 | 89 | # Initiate graceful closing of the AMQP broker connection 90 | self._conn.close() 91 | return 92 | 93 | def _connection_closed_cb(self): 94 | print "AMQP broker connection closed; close-info: %s" % ( 95 | self._conn.close_info,) 96 | self._conn = None 97 | return 98 | 99 | 100 | def main(): 101 | waiter = gevent_event.AsyncResult() 102 | 103 | HaighaGeventHello(waiter.set) 104 | 105 | print "Waiting for I/O to complete..." 106 | waiter.wait() 107 | 108 | print "Done!" 109 | return 110 | 111 | 112 | 113 | if __name__ == '__main__': 114 | logging.basicConfig() 115 | main() 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | 124 | 125 | -------------------------------------------------------------------------------- /examples/rpc_client.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | ''' 3 | demostrate how to write a rpc client 4 | ''' 5 | import sys, os, uuid, time 6 | sys.path.append(os.path.abspath("..")) 7 | 8 | from haigha.connection import Connection 9 | from haigha.message import Message 10 | 11 | class FibonacciRpcClient(object): 12 | def __init__(self): 13 | self.connection = Connection(host='localhost', heartbeat=None, debug=True) 14 | 15 | self.channel = self.connection.channel() 16 | 17 | result = self.channel.queue.declare(exclusive=True) 18 | self.callback_queue = result[0] 19 | print("callback_queue:", self.callback_queue) 20 | 21 | self.channel.basic.consume(self.callback_queue, self.on_response, no_ack=True) 22 | 23 | def on_response(self, msg): 24 | if msg.properties["correlation_id"] == self.corr_id: 25 | self.response = msg.body 26 | 27 | def call(self, n): 28 | self.response = None 29 | self.corr_id = str(uuid.uuid4()) 30 | msg = Message(str(n), reply_to=self.callback_queue, correlation_id=self.corr_id) 31 | self.channel.basic.publish(msg, '', 'rpc_queue') 32 | while self.response is None: 33 | self.connection.read_frames() 34 | return int(self.response) 35 | 36 | fibonacci_rpc = FibonacciRpcClient() 37 | 38 | print " [x] Requesting fib(30)" 39 | response = fibonacci_rpc.call(30) 40 | print " [.] Got %r" % (response,) 41 | 42 | -------------------------------------------------------------------------------- /examples/rpc_server.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | ''' 3 | demostrate how to write a rpc server 4 | ''' 5 | import sys, os, uuid, time 6 | sys.path.append(os.path.abspath("..")) 7 | 8 | from haigha.connection import Connection 9 | from haigha.message import Message 10 | 11 | connection = Connection(host='localhost', heartbeat=None, debug=True) 12 | channel = connection.channel() 13 | channel.queue.declare(queue='rpc_queue', auto_delete=False) 14 | 15 | def fib(n): 16 | if n == 0: 17 | return 0 18 | elif n == 1: 19 | return 1 20 | else: 21 | return fib(n-1) + fib(n-2) 22 | 23 | def on_request(msg): 24 | n = int(msg.body) 25 | 26 | print " [.] fib(%s)" % (n,) 27 | result = fib(n) 28 | 29 | reply_to = msg.properties["reply_to"] 30 | correlation_id = msg.properties["correlation_id"] 31 | resp = Message(str(result), correlation_id=correlation_id) 32 | channel.basic.publish(resp,'',reply_to) 33 | 34 | delivery_info = msg.delivery_info 35 | channel.basic.ack(delivery_info["delivery_tag"]) 36 | 37 | channel.basic.qos(prefetch_count=1) 38 | channel.basic.consume('rpc_queue', on_request, no_ack=False) 39 | 40 | print " [x] Awaiting RPC requests" 41 | while not channel.closed: 42 | connection.read_frames() 43 | 44 | -------------------------------------------------------------------------------- /examples/synchronous: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #-*- coding:utf-8 -*- 3 | 4 | import sys, os 5 | sys.path.append(os.path.abspath(".")) 6 | sys.path.append(os.path.abspath("..")) 7 | 8 | import logging 9 | import random 10 | import socket 11 | from optparse import OptionParser 12 | 13 | from haigha.connection import Connection 14 | from haigha.message import Message 15 | 16 | parser = OptionParser( 17 | usage='Usage: synchronous_test [options]' 18 | ) 19 | parser.add_option('--user', default='guest', type='string') 20 | parser.add_option('--pass', default='guest', dest='password', type='string') 21 | parser.add_option('--vhost', default='/', type='string') 22 | parser.add_option('--host', default='localhost', type='string') 23 | parser.add_option('--debug', default=0, action='count') 24 | 25 | (options,args) = parser.parse_args() 26 | 27 | debug = options.debug 28 | level = logging.DEBUG if debug else logging.INFO 29 | 30 | # Setup logging 31 | logging.basicConfig(level=level, format="[%(levelname)s %(asctime)s] %(message)s" ) 32 | logger = logging.getLogger('haigha') 33 | 34 | sock_opts = { 35 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 36 | } 37 | connection = Connection(logger=logger, debug=debug, 38 | user=options.user, password=options.password, 39 | vhost=options.vhost, host=options.host, 40 | heartbeat=None, 41 | sock_opts=sock_opts, 42 | transport='socket') 43 | 44 | ch = connection.channel() 45 | ch.exchange.declare('foo', 'direct') 46 | ch.queue.declare('bar') 47 | ch.queue.bind('bar', 'foo', 'route') 48 | ch.basic.publish(Message('hello world'), 'foo', 'route') 49 | print 'GET:', ch.basic.get('bar') 50 | 51 | ch.basic.publish(Message('hello world'), 'foo', 'route') 52 | ch.basic.publish(Message('hello world'), 'foo', 'route') 53 | print 'PURGE:', ch.queue.purge('bar') 54 | 55 | ch.basic.publish(Message('hello world'), 'foo', 'route') 56 | ch.basic.publish(Message('hello world'), 'foo', 'route') 57 | print 'DELETED:', ch.queue.delete('bar') 58 | -------------------------------------------------------------------------------- /examples/synchronous_gevent: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #-*- coding:utf-8 -*- 3 | 4 | import sys, os 5 | sys.path.append(os.path.abspath(".")) 6 | sys.path.append(os.path.abspath("..")) 7 | 8 | import logging 9 | import random 10 | import socket 11 | from optparse import OptionParser 12 | 13 | from haigha.connection import Connection 14 | from haigha.message import Message 15 | 16 | parser = OptionParser( 17 | usage='Usage: synchronous_test [options]' 18 | ) 19 | parser.add_option('--user', default='guest', type='string') 20 | parser.add_option('--pass', default='guest', dest='password', type='string') 21 | parser.add_option('--vhost', default='/', type='string') 22 | parser.add_option('--host', default='localhost', type='string') 23 | parser.add_option('--debug', default=0, action='count') 24 | 25 | (options,args) = parser.parse_args() 26 | 27 | debug = options.debug 28 | level = logging.DEBUG if debug else logging.INFO 29 | 30 | # Setup logging 31 | logging.basicConfig(level=level, format="[%(levelname)s %(asctime)s] %(message)s" ) 32 | logger = logging.getLogger('haigha') 33 | 34 | sock_opts = { 35 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 36 | } 37 | connection = Connection(logger=logger, debug=debug, 38 | user=options.user, password=options.password, 39 | vhost=options.vhost, host=options.host, 40 | heartbeat=None, 41 | sock_opts=sock_opts, 42 | transport='gevent', 43 | synchronous=True) 44 | 45 | ch = connection.channel() 46 | ch.exchange.declare('foo', 'direct') 47 | ch.queue.declare('bar') 48 | ch.queue.bind('bar', 'foo', 'route') 49 | ch.basic.publish(Message('hello world'), 'foo', 'route') 50 | print 'GET:', ch.basic.get('bar') 51 | 52 | ch.basic.publish(Message('hello world'), 'foo', 'route') 53 | ch.basic.publish(Message('hello world'), 'foo', 'route') 54 | print 'PURGE:', ch.queue.purge('bar') 55 | 56 | ch.basic.publish(Message('hello world'), 'foo', 'route') 57 | ch.basic.publish(Message('hello world'), 'foo', 'route') 58 | print 'DELETED:', ch.queue.delete('bar') 59 | 60 | -------------------------------------------------------------------------------- /haigha/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | __version__ = "0.9.0" 8 | -------------------------------------------------------------------------------- /haigha/channel_pool.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from collections import deque 8 | 9 | 10 | class ChannelPool(object): 11 | 12 | ''' 13 | Manages a pool of channels for transaction-based publishing. This allows a 14 | client to use as many channels as are necessary to publish while not 15 | creating a backlog of transactions that slows throughput and consumes 16 | memory. 17 | 18 | The pool can accept an optional `size` argument in the ctor, which caps 19 | the number of channels which the pool will allocate. If no channels are 20 | available on `publish()`, the message will be locally queued and sent as 21 | soon as a channel is available. It is recommended that you use the pool 22 | with a max size, as each channel consumes memory on the broker and it is 23 | possible to exercise memory limit protection seems on the broker due to 24 | number of channels. 25 | ''' 26 | 27 | def __init__(self, connection, size=None): 28 | '''Initialize the channel on a connection.''' 29 | self._connection = connection 30 | self._free_channels = set() 31 | self._size = size 32 | self._queue = deque() 33 | self._channels = 0 34 | 35 | def publish(self, *args, **kwargs): 36 | ''' 37 | Publish a message. Caller can supply an optional callback which will 38 | be fired when the transaction is committed. Tries very hard to avoid 39 | closed and inactive channels, but a ChannelError or ConnectionError 40 | may still be raised. 41 | ''' 42 | user_cb = kwargs.pop('cb', None) 43 | 44 | # If the first channel we grab is inactive, continue fetching until 45 | # we get an active channel, then put the inactive channels back in 46 | # the pool. Try to keep the overhead to a minimum. 47 | channel = self._get_channel() 48 | 49 | if channel and not channel.active: 50 | inactive_channels = set() 51 | while channel and not channel.active: 52 | inactive_channels.add(channel) 53 | channel = self._get_channel() 54 | self._free_channels.update(inactive_channels) 55 | 56 | # When the transaction is committed, add the channel back to the pool 57 | # and call any user-defined callbacks. If there is anything in queue, 58 | # pop it and call back to publish(). Only do so if the channel is 59 | # still active though, because otherwise the message will end up at 60 | # the back of the queue, breaking the original order. 61 | def committed(): 62 | self._free_channels.add(channel) 63 | if channel.active and not channel.closed: 64 | self._process_queue() 65 | if user_cb is not None: 66 | user_cb() 67 | 68 | if channel: 69 | channel.publish_synchronous(*args, cb=committed, **kwargs) 70 | else: 71 | kwargs['cb'] = user_cb 72 | self._queue.append((args, kwargs)) 73 | 74 | def _process_queue(self): 75 | ''' 76 | If there are any message in the queue, process one of them. 77 | ''' 78 | if len(self._queue): 79 | args, kwargs = self._queue.popleft() 80 | self.publish(*args, **kwargs) 81 | 82 | def _get_channel(self): 83 | ''' 84 | Fetch a channel from the pool. Will return a new one if necessary. If 85 | a channel in the free pool is closed, will remove it. Will return None 86 | if we hit the cap. Will clean up any channels that were published to 87 | but closed due to error. 88 | ''' 89 | while len(self._free_channels): 90 | rval = self._free_channels.pop() 91 | if not rval.closed: 92 | return rval 93 | # don't adjust _channels value because the callback will do that 94 | # and we don't want to double count it. 95 | 96 | if not self._size or self._channels < self._size: 97 | rval = self._connection.channel() 98 | self._channels += 1 99 | rval.add_close_listener(self._channel_closed_cb) 100 | return rval 101 | 102 | def _channel_closed_cb(self, channel): 103 | ''' 104 | Callback when channel closes. 105 | ''' 106 | self._channels -= 1 107 | self._process_queue() 108 | -------------------------------------------------------------------------------- /haigha/classes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/haigha/classes/__init__.py -------------------------------------------------------------------------------- /haigha/classes/channel_class.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.classes.protocol_class import ProtocolClass 8 | from haigha.frames.method_frame import MethodFrame 9 | from haigha.writer import Writer 10 | 11 | 12 | class ChannelClass(ProtocolClass): 13 | 14 | ''' 15 | Implements the AMQP Channel class 16 | ''' 17 | 18 | CLASS_ID = 20 19 | 20 | # Channel method ids for error-recovery code in Channel 21 | CLOSE_METHOD_ID = 40 22 | CLOSE_OK_METHOD_ID = 41 23 | 24 | def __init__(self, *args, **kwargs): 25 | super(ChannelClass, self).__init__(*args, **kwargs) 26 | self.dispatch_map = { 27 | 11: self._recv_open_ok, 28 | 20: self._recv_flow, 29 | 21: self._recv_flow_ok, 30 | 40: self._recv_close, 31 | 41: self._recv_close_ok, 32 | } 33 | self._flow_control_cb = None 34 | 35 | @property 36 | def name(self): 37 | return 'channel' 38 | 39 | def set_flow_cb(self, cb): 40 | ''' 41 | Set a callback that will be called when the state of flow control has 42 | changed. The caller should use closures if they need to receive a 43 | handle to the channel on which flow control changes. 44 | ''' 45 | self._flow_control_cb = cb 46 | 47 | def open(self): 48 | ''' 49 | Open the channel for communication. 50 | ''' 51 | args = Writer() 52 | args.write_shortstr('') 53 | self.send_frame(MethodFrame(self.channel_id, 20, 10, args)) 54 | self.channel.add_synchronous_cb(self._recv_open_ok) 55 | 56 | def _recv_open_ok(self, method_frame): 57 | ''' 58 | Channel is opened. 59 | ''' 60 | self.channel._notify_open_listeners() 61 | 62 | def activate(self): 63 | ''' 64 | Activate this channel (disable flow control). 65 | ''' 66 | if not self.channel.active: 67 | self._send_flow(True) 68 | 69 | def deactivate(self): 70 | ''' 71 | Deactivate this channel (enable flow control). 72 | ''' 73 | if self.channel.active: 74 | self._send_flow(False) 75 | 76 | def _send_flow(self, active): 77 | ''' 78 | Send a flow control command. 79 | ''' 80 | args = Writer() 81 | args.write_bit(active) 82 | self.send_frame(MethodFrame(self.channel_id, 20, 20, args)) 83 | self.channel.add_synchronous_cb(self._recv_flow_ok) 84 | 85 | def _recv_flow(self, method_frame): 86 | ''' 87 | Receive a flow control command from the broker 88 | ''' 89 | self.channel._active = method_frame.args.read_bit() 90 | 91 | args = Writer() 92 | args.write_bit(self.channel.active) 93 | self.send_frame(MethodFrame(self.channel_id, 20, 21, args)) 94 | 95 | if self._flow_control_cb is not None: 96 | self._flow_control_cb() 97 | 98 | def _recv_flow_ok(self, method_frame): 99 | ''' 100 | Receive a flow control ack from the broker. 101 | ''' 102 | self.channel._active = method_frame.args.read_bit() 103 | if self._flow_control_cb is not None: 104 | self._flow_control_cb() 105 | 106 | def close(self, reply_code=0, reply_text='', class_id=0, method_id=0): 107 | ''' 108 | Close this channel. Caller has the option of specifying the reason for 109 | closure and the class and method ids of the current frame in which an 110 | error occurred. If in the event of an exception, the channel will be 111 | marked as immediately closed. If channel is already closed, call is 112 | ignored. 113 | ''' 114 | if not getattr(self, 'channel', None) or self.channel._closed: 115 | return 116 | 117 | self.channel._close_info = { 118 | 'reply_code': reply_code, 119 | 'reply_text': reply_text, 120 | 'class_id': class_id, 121 | 'method_id': method_id 122 | } 123 | 124 | # exceptions here likely due to race condition as connection is closing 125 | # cap the reply_text we send because it may be arbitrarily long 126 | try: 127 | args = Writer() 128 | args.write_short(reply_code) 129 | args.write_shortstr(reply_text[:255]) 130 | args.write_short(class_id) 131 | args.write_short(method_id) 132 | self.send_frame(MethodFrame(self.channel_id, 20, 40, args)) 133 | 134 | self.channel.add_synchronous_cb(self._recv_close_ok) 135 | finally: 136 | # Immediately set the closed flag so no more frames can be sent 137 | # NOTE: in synchronous mode, by the time this is called we will 138 | # have already run self.channel._closed_cb and so the channel 139 | # reference is gone. 140 | if self.channel: 141 | self.channel._closed = True 142 | 143 | def _recv_close(self, method_frame): 144 | ''' 145 | Receive a close command from the broker. 146 | ''' 147 | self.channel._close_info = { 148 | 'reply_code': method_frame.args.read_short(), 149 | 'reply_text': method_frame.args.read_shortstr(), 150 | 'class_id': method_frame.args.read_short(), 151 | 'method_id': method_frame.args.read_short() 152 | } 153 | 154 | self.channel._closed = True 155 | self.channel._closed_cb( 156 | final_frame=MethodFrame(self.channel_id, 20, 41)) 157 | 158 | def _recv_close_ok(self, method_frame): 159 | ''' 160 | Receive a close ack from the broker. 161 | ''' 162 | self.channel._closed = True 163 | self.channel._closed_cb() 164 | -------------------------------------------------------------------------------- /haigha/classes/exchange_class.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from collections import deque 8 | 9 | from haigha.writer import Writer 10 | from haigha.frames.method_frame import MethodFrame 11 | from haigha.classes.protocol_class import ProtocolClass 12 | 13 | 14 | class ExchangeClass(ProtocolClass): 15 | 16 | ''' 17 | Implements the AMQP Exchange class 18 | ''' 19 | 20 | def __init__(self, *args, **kwargs): 21 | super(ExchangeClass, self).__init__(*args, **kwargs) 22 | self.dispatch_map = { 23 | 11: self._recv_declare_ok, 24 | 21: self._recv_delete_ok, 25 | } 26 | 27 | self._declare_cb = deque() 28 | self._delete_cb = deque() 29 | 30 | @property 31 | def name(self): 32 | return 'exchange' 33 | 34 | def _cleanup(self): 35 | ''' 36 | Cleanup local data. 37 | ''' 38 | self._declare_cb = None 39 | self._delete_cb = None 40 | super(ExchangeClass, self)._cleanup() 41 | 42 | def declare(self, exchange, type, passive=False, durable=False, 43 | nowait=True, arguments=None, ticket=None, cb=None): 44 | """ 45 | Declare the exchange. 46 | 47 | exchange - The name of the exchange to declare 48 | type - One of 49 | """ 50 | nowait = nowait and self.allow_nowait() and not cb 51 | 52 | args = Writer() 53 | args.write_short(ticket or self.default_ticket).\ 54 | write_shortstr(exchange).\ 55 | write_shortstr(type).\ 56 | write_bits(passive, durable, False, False, nowait).\ 57 | write_table(arguments or {}) 58 | self.send_frame(MethodFrame(self.channel_id, 40, 10, args)) 59 | 60 | if not nowait: 61 | self._declare_cb.append(cb) 62 | self.channel.add_synchronous_cb(self._recv_declare_ok) 63 | 64 | def _recv_declare_ok(self, _method_frame): 65 | ''' 66 | Confirmation that exchange was declared. 67 | ''' 68 | cb = self._declare_cb.popleft() 69 | if cb: 70 | cb() 71 | 72 | def delete(self, exchange, if_unused=False, nowait=True, ticket=None, 73 | cb=None): 74 | ''' 75 | Delete an exchange. 76 | ''' 77 | nowait = nowait and self.allow_nowait() and not cb 78 | 79 | args = Writer() 80 | args.write_short(ticket or self.default_ticket).\ 81 | write_shortstr(exchange).\ 82 | write_bits(if_unused, nowait) 83 | self.send_frame(MethodFrame(self.channel_id, 40, 20, args)) 84 | 85 | if not nowait: 86 | self._delete_cb.append(cb) 87 | self.channel.add_synchronous_cb(self._recv_delete_ok) 88 | 89 | def _recv_delete_ok(self, _method_frame): 90 | ''' 91 | Confirmation that exchange was deleted. 92 | ''' 93 | cb = self._delete_cb.popleft() 94 | if cb: 95 | cb() 96 | -------------------------------------------------------------------------------- /haigha/classes/protocol_class.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | 8 | class ProtocolClass(object): 9 | 10 | ''' 11 | The base class of all protocol classes. 12 | ''' 13 | 14 | class ProtocolError(Exception): 15 | pass 16 | 17 | class InvalidMethod(ProtocolError): 18 | pass 19 | 20 | class FrameUnderflow(ProtocolError): 21 | pass 22 | 23 | dispatch_map = {} 24 | 25 | def __init__(self, channel): 26 | ''' 27 | Construct this protocol class on a channel. 28 | ''' 29 | # Cache the channel id so that cleanup can remove the circular channel 30 | # reference but id is still accessible (it's useful!) 31 | self._channel = channel 32 | self._channel_id = channel.channel_id 33 | 34 | @property 35 | def channel(self): 36 | return self._channel 37 | 38 | @property 39 | def channel_id(self): 40 | return self._channel_id 41 | 42 | @property 43 | def logger(self): 44 | return self._channel.logger 45 | 46 | @property 47 | def default_ticket(self): 48 | return 0 49 | 50 | @property 51 | def name(self): 52 | '''The name given this in the protocol, i.e. 'basic', 'tx', etc''' 53 | raise NotImplementedError('must provide a name for %s' % (self)) 54 | 55 | def allow_nowait(self): 56 | ''' 57 | Return True if the transport or channel allows nowait, 58 | False otherwise. 59 | ''' 60 | return not self._channel.synchronous 61 | 62 | def _cleanup(self): 63 | ''' 64 | "Private" call from Channel when it's shutting down so that local 65 | data can be cleaned up and references closed out. It's strongly 66 | recommended that subclasses call this /after/ doing their own cleanup . 67 | Note that this removes reference to both the channel and the dispatch 68 | map. 69 | ''' 70 | self._channel = None 71 | self.dispatch_map = None 72 | 73 | def dispatch(self, method_frame): 74 | ''' 75 | Dispatch a method for this protocol. 76 | ''' 77 | method = self.dispatch_map.get(method_frame.method_id) 78 | if method: 79 | callback = self.channel.clear_synchronous_cb(method) 80 | callback(method_frame) 81 | else: 82 | raise self.InvalidMethod( 83 | "no method is registered with id: %d" % method_frame.method_id) 84 | 85 | def send_frame(self, frame): 86 | ''' 87 | Send a frame 88 | ''' 89 | self.channel.send_frame(frame) 90 | -------------------------------------------------------------------------------- /haigha/classes/queue_class.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.writer import Writer 8 | from haigha.frames.method_frame import MethodFrame 9 | from haigha.classes.protocol_class import ProtocolClass 10 | 11 | from collections import deque 12 | 13 | 14 | class QueueClass(ProtocolClass): 15 | 16 | ''' 17 | Implements the AMQP Queue class 18 | ''' 19 | 20 | def __init__(self, *args, **kwargs): 21 | super(QueueClass, self).__init__(*args, **kwargs) 22 | self.dispatch_map = { 23 | 11: self._recv_declare_ok, 24 | 21: self._recv_bind_ok, 25 | 31: self._recv_purge_ok, 26 | 41: self._recv_delete_ok, 27 | 51: self._recv_unbind_ok, 28 | } 29 | 30 | self._declare_cb = deque() 31 | self._bind_cb = deque() 32 | self._unbind_cb = deque() 33 | self._delete_cb = deque() 34 | self._purge_cb = deque() 35 | 36 | @property 37 | def name(self): 38 | return 'queue' 39 | 40 | def _cleanup(self): 41 | ''' 42 | Cleanup all the local data. 43 | ''' 44 | self._declare_cb = None 45 | self._bind_cb = None 46 | self._unbind_cb = None 47 | self._delete_cb = None 48 | self._purge_cb = None 49 | super(QueueClass, self)._cleanup() 50 | 51 | def declare(self, queue='', passive=False, durable=False, 52 | exclusive=False, auto_delete=True, nowait=True, 53 | arguments={}, ticket=None, cb=None): 54 | ''' 55 | Declare a queue. By default is asynchronoous but will be synchronous 56 | if nowait=False or a callback is defined. In synchronous mode, 57 | returns (message_count, consumer_count) 58 | 59 | queue - The name of the queue 60 | cb - An optional method which will be called with 61 | (queue_name, msg_count, consumer_count) if nowait=False 62 | ''' 63 | nowait = nowait and self.allow_nowait() and not cb 64 | 65 | args = Writer() 66 | args.write_short(ticket or self.default_ticket).\ 67 | write_shortstr(queue).\ 68 | write_bits(passive, durable, exclusive, auto_delete, nowait).\ 69 | write_table(arguments) 70 | self.send_frame(MethodFrame(self.channel_id, 50, 10, args)) 71 | 72 | if not nowait: 73 | self._declare_cb.append(cb) 74 | return self.channel.add_synchronous_cb(self._recv_declare_ok) 75 | 76 | def _recv_declare_ok(self, method_frame): 77 | queue = method_frame.args.read_shortstr() 78 | message_count = method_frame.args.read_long() 79 | consumer_count = method_frame.args.read_long() 80 | 81 | cb = self._declare_cb.popleft() 82 | if cb: 83 | cb(queue, message_count, consumer_count) 84 | return queue, message_count, consumer_count 85 | 86 | def bind(self, queue, exchange, routing_key='', nowait=True, arguments={}, 87 | ticket=None, cb=None): 88 | ''' 89 | bind to a queue. 90 | ''' 91 | nowait = nowait and self.allow_nowait() and not cb 92 | 93 | args = Writer() 94 | args.write_short(ticket or self.default_ticket).\ 95 | write_shortstr(queue).\ 96 | write_shortstr(exchange).\ 97 | write_shortstr(routing_key).\ 98 | write_bit(nowait).\ 99 | write_table(arguments) 100 | self.send_frame(MethodFrame(self.channel_id, 50, 20, args)) 101 | 102 | if not nowait: 103 | self._bind_cb.append(cb) 104 | self.channel.add_synchronous_cb(self._recv_bind_ok) 105 | 106 | def _recv_bind_ok(self, _method_frame): 107 | # No arguments defined. 108 | cb = self._bind_cb.popleft() 109 | if cb: 110 | cb() 111 | 112 | def unbind(self, queue, exchange, routing_key='', arguments={}, 113 | ticket=None, cb=None): 114 | ''' 115 | Unbind a queue from an exchange. This is always synchronous. 116 | ''' 117 | args = Writer() 118 | args.write_short(ticket or self.default_ticket).\ 119 | write_shortstr(queue).\ 120 | write_shortstr(exchange).\ 121 | write_shortstr(routing_key).\ 122 | write_table(arguments) 123 | self.send_frame(MethodFrame(self.channel_id, 50, 50, args)) 124 | 125 | self._unbind_cb.append(cb) 126 | self.channel.add_synchronous_cb(self._recv_unbind_ok) 127 | 128 | def _recv_unbind_ok(self, _method_frame): 129 | # No arguments defined 130 | cb = self._unbind_cb.popleft() 131 | if cb: 132 | cb() 133 | 134 | def purge(self, queue, nowait=True, ticket=None, cb=None): 135 | ''' 136 | Purge all messages in a queue. 137 | ''' 138 | nowait = nowait and self.allow_nowait() and not cb 139 | 140 | args = Writer() 141 | args.write_short(ticket or self.default_ticket).\ 142 | write_shortstr(queue).\ 143 | write_bit(nowait) 144 | self.send_frame(MethodFrame(self.channel_id, 50, 30, args)) 145 | 146 | if not nowait: 147 | self._purge_cb.append(cb) 148 | return self.channel.add_synchronous_cb(self._recv_purge_ok) 149 | 150 | def _recv_purge_ok(self, method_frame): 151 | message_count = method_frame.args.read_long() 152 | cb = self._purge_cb.popleft() 153 | if cb: 154 | cb(message_count) 155 | return message_count 156 | 157 | def delete(self, queue, if_unused=False, if_empty=False, nowait=True, 158 | ticket=None, cb=None): 159 | ''' 160 | queue delete. 161 | ''' 162 | nowait = nowait and self.allow_nowait() and not cb 163 | 164 | args = Writer() 165 | args.write_short(ticket or self.default_ticket).\ 166 | write_shortstr(queue).\ 167 | write_bits(if_unused, if_empty, nowait) 168 | self.send_frame(MethodFrame(self.channel_id, 50, 40, args)) 169 | 170 | if not nowait: 171 | self._delete_cb.append(cb) 172 | return self.channel.add_synchronous_cb(self._recv_delete_ok) 173 | 174 | def _recv_delete_ok(self, method_frame): 175 | message_count = method_frame.args.read_long() 176 | cb = self._delete_cb.popleft() 177 | if cb: 178 | cb(message_count) 179 | return message_count 180 | -------------------------------------------------------------------------------- /haigha/classes/transaction_class.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.frames.method_frame import MethodFrame 8 | from haigha.classes.protocol_class import ProtocolClass 9 | 10 | from collections import deque 11 | 12 | 13 | class TransactionClass(ProtocolClass): 14 | 15 | ''' 16 | Implements the AMQP Transaction class 17 | ''' 18 | 19 | class TransactionsNotEnabled(ProtocolClass.ProtocolError): 20 | 21 | '''Tried to use transactions without enabling them.''' 22 | 23 | def __init__(self, *args, **kwargs): 24 | super(TransactionClass, self).__init__(*args, **kwargs) 25 | self.dispatch_map = { 26 | 11: self._recv_select_ok, 27 | 21: self._recv_commit_ok, 28 | 31: self._recv_rollback_ok, 29 | } 30 | 31 | self._enabled = False 32 | self._select_cb = deque() 33 | self._commit_cb = deque() 34 | self._rollback_cb = deque() 35 | 36 | @property 37 | def name(self): 38 | return 'tx' 39 | 40 | @property 41 | def enabled(self): 42 | '''Get whether transactions have been enabled.''' 43 | return self._enabled 44 | 45 | def _cleanup(self): 46 | ''' 47 | Cleanup all the local data. 48 | ''' 49 | self._select_cb = None 50 | self._commit_cb = None 51 | self._rollback_cb = None 52 | super(TransactionClass, self)._cleanup() 53 | 54 | def select(self, cb=None): 55 | ''' 56 | Set this channel to use transactions. 57 | ''' 58 | if not self._enabled: 59 | self._enabled = True 60 | self.send_frame(MethodFrame(self.channel_id, 90, 10)) 61 | self._select_cb.append(cb) 62 | self.channel.add_synchronous_cb(self._recv_select_ok) 63 | 64 | def _recv_select_ok(self, _method_frame): 65 | cb = self._select_cb.popleft() 66 | if cb: 67 | cb() 68 | 69 | def commit(self, cb=None): 70 | ''' 71 | Commit the current transaction. Caller can specify a callback to use 72 | when the transaction is committed. 73 | ''' 74 | # Could call select() but spec 1.9.2.3 says to raise an exception 75 | if not self.enabled: 76 | raise self.TransactionsNotEnabled() 77 | 78 | self.send_frame(MethodFrame(self.channel_id, 90, 20)) 79 | self._commit_cb.append(cb) 80 | self.channel.add_synchronous_cb(self._recv_commit_ok) 81 | 82 | def _recv_commit_ok(self, _method_frame): 83 | cb = self._commit_cb.popleft() 84 | if cb: 85 | cb() 86 | 87 | def rollback(self, cb=None): 88 | ''' 89 | Abandon all message publications and acks in the current transaction. 90 | Caller can specify a callback to use when the transaction has been 91 | aborted. 92 | ''' 93 | # Could call select() but spec 1.9.2.5 says to raise an exception 94 | if not self.enabled: 95 | raise self.TransactionsNotEnabled() 96 | 97 | self.send_frame(MethodFrame(self.channel_id, 90, 30)) 98 | self._rollback_cb.append(cb) 99 | self.channel.add_synchronous_cb(self._recv_rollback_ok) 100 | 101 | def _recv_rollback_ok(self, _method_frame): 102 | cb = self._rollback_cb.popleft() 103 | if cb: 104 | cb() 105 | -------------------------------------------------------------------------------- /haigha/connections/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | __all__ = ['rabbit_connection'] 8 | -------------------------------------------------------------------------------- /haigha/exceptions.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | 8 | class ConnectionError(Exception): 9 | 10 | '''Base class for all connection errors.''' 11 | 12 | 13 | class ConnectionClosed(ConnectionError): 14 | 15 | '''The connection is closed. Fatal.''' 16 | 17 | 18 | class ChannelError(Exception): 19 | 20 | '''Base class for all channel errors.''' 21 | 22 | 23 | class ChannelClosed(ChannelError): 24 | 25 | '''The channel is closed. Fatal.''' 26 | -------------------------------------------------------------------------------- /haigha/frames/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/haigha/frames/__init__.py -------------------------------------------------------------------------------- /haigha/frames/content_frame.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.writer import Writer 8 | from haigha.frames.frame import Frame 9 | 10 | 11 | class ContentFrame(Frame): 12 | 13 | ''' 14 | Frame for reading in content. 15 | ''' 16 | 17 | @classmethod 18 | def type(cls): 19 | return 3 20 | 21 | @property 22 | def payload(self): 23 | return self._payload 24 | 25 | @classmethod 26 | def parse(self, channel_id, payload): 27 | return ContentFrame(channel_id, payload) 28 | 29 | @classmethod 30 | def create_frames(self, channel_id, buf, frame_max): 31 | ''' 32 | A generator which will create frames from a buffer given a max 33 | frame size. 34 | ''' 35 | size = frame_max - 8 # 8 bytes overhead for frame header and footer 36 | offset = 0 37 | while True: 38 | payload = buf[offset:(offset + size)] 39 | if len(payload) == 0: 40 | break 41 | offset += size 42 | 43 | yield ContentFrame(channel_id, payload) 44 | if offset >= len(buf): 45 | break 46 | 47 | def __init__(self, channel_id, payload): 48 | Frame.__init__(self, channel_id) 49 | self._payload = payload 50 | 51 | def __str__(self): 52 | if isinstance(self._payload, str): 53 | payload = ''.join(['\\x%s' % (c.encode('hex')) 54 | for c in self._payload]) 55 | else: 56 | payload = str(self._payload) 57 | 58 | return "%s[channel: %d, payload: %s]" % ( 59 | self.__class__.__name__, self.channel_id, payload) 60 | 61 | def write_frame(self, buf): 62 | ''' 63 | Write the frame into an existing buffer. 64 | ''' 65 | writer = Writer(buf) 66 | 67 | writer.write_octet(self.type()).\ 68 | write_short(self.channel_id).\ 69 | write_long(len(self._payload)).\ 70 | write(self._payload).\ 71 | write_octet(0xce) 72 | 73 | 74 | ContentFrame.register() 75 | -------------------------------------------------------------------------------- /haigha/frames/frame.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | import struct 8 | import sys 9 | from collections import deque 10 | from haigha.reader import Reader 11 | 12 | 13 | class Frame(object): 14 | 15 | ''' 16 | Base class for a frame. 17 | ''' 18 | 19 | # Exceptions 20 | class FrameError(Exception): 21 | 22 | '''Base class for all frame errors''' 23 | class FormatError(FrameError): 24 | 25 | '''The frame was mal-formed.''' 26 | class InvalidFrameType(FrameError): 27 | 28 | '''The frame type is unknown.''' 29 | 30 | # Class data 31 | _frame_type_map = {} 32 | 33 | # Class methods 34 | @classmethod 35 | def register(cls): 36 | ''' 37 | Register a frame type. 38 | ''' 39 | cls._frame_type_map[cls.type()] = cls 40 | 41 | @classmethod 42 | def type(self): 43 | ''' 44 | Fetch the type of this frame. Should be an octet. 45 | ''' 46 | raise NotImplementedError() 47 | 48 | @classmethod 49 | def read_frames(cls, reader): 50 | ''' 51 | Read one or more frames from an IO stream. Buffer must support file 52 | object interface. 53 | 54 | After reading, caller will need to check if there are bytes remaining 55 | in the stream. If there are, then that implies that there is one or 56 | more incomplete frames and more data needs to be read. The position 57 | of the cursor in the frame stream will mark the point at which the 58 | last good frame was read. If the caller is expecting a sequence of 59 | frames and only received a part of that sequence, they are responsible 60 | for buffering those frames until the rest of the frames in the sequence 61 | have arrived. 62 | ''' 63 | rval = deque() 64 | 65 | while True: 66 | frame_start_pos = reader.tell() 67 | try: 68 | frame = Frame._read_frame(reader) 69 | except Reader.BufferUnderflow: 70 | # No more data in the stream 71 | frame = None 72 | except Reader.ReaderError as e: 73 | # Some other format error 74 | raise Frame.FormatError, str(e), sys.exc_info()[-1] 75 | except struct.error as e: 76 | raise Frame.FormatError, str(e), sys.exc_info()[-1] 77 | 78 | if frame is None: 79 | reader.seek(frame_start_pos) 80 | break 81 | 82 | rval.append(frame) 83 | 84 | return rval 85 | 86 | @classmethod 87 | def _read_frame(cls, reader): 88 | ''' 89 | Read a single frame from a Reader. Will return None if there is an 90 | incomplete frame in the stream. 91 | 92 | Raise MissingFooter if there's a problem reading the footer byte. 93 | ''' 94 | frame_type = reader.read_octet() 95 | channel_id = reader.read_short() 96 | size = reader.read_long() 97 | 98 | payload = Reader(reader, reader.tell(), size) 99 | 100 | # Seek to end of payload 101 | reader.seek(size, 1) 102 | 103 | ch = reader.read_octet() # footer 104 | if ch != 0xce: 105 | raise Frame.FormatError( 106 | 'Framing error, unexpected byte: %x. frame type %x. channel %d, payload size %d', 107 | ch, frame_type, channel_id, size) 108 | 109 | frame_class = cls._frame_type_map.get(frame_type) 110 | if not frame_class: 111 | raise Frame.InvalidFrameType("Unknown frame type %x", frame_type) 112 | return frame_class.parse(channel_id, payload) 113 | 114 | # Instance methods 115 | def __init__(self, channel_id=-1): 116 | self._channel_id = channel_id 117 | 118 | @classmethod 119 | def parse(cls, channel_id, payload): 120 | ''' 121 | Subclasses need to implement parsing of their frames. Should return 122 | a new instance of their type. 123 | ''' 124 | raise NotImplementedError() 125 | 126 | @property 127 | def channel_id(self): 128 | return self._channel_id 129 | 130 | def __str__(self): 131 | return "%s[channel: %d]" % (self.__class__.__name__, self.channel_id) 132 | 133 | def __repr__(self): 134 | # Have to actually call the method rather than __repr__==__str__ 135 | # because subclasses overload __str__ 136 | return str(self) 137 | 138 | def write_frame(self, stream): 139 | ''' 140 | Write this frame. 141 | ''' 142 | raise NotImplementedError() 143 | -------------------------------------------------------------------------------- /haigha/frames/header_frame.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from collections import deque 8 | 9 | from haigha.writer import Writer 10 | from haigha.reader import Reader 11 | from haigha.frames.frame import Frame 12 | 13 | 14 | class HeaderFrame(Frame): 15 | 16 | ''' 17 | Header frame for content. 18 | ''' 19 | PROPERTIES = [ 20 | ('content_type', 'shortstr', Reader.read_shortstr, 21 | Writer.write_shortstr, 1 << 15), 22 | ('content_encoding', 'shortstr', Reader.read_shortstr, 23 | Writer.write_shortstr, 1 << 14), 24 | ('application_headers', 'table', 25 | Reader.read_table, Writer.write_table, 1 << 13), 26 | ('delivery_mode', 'octet', Reader.read_octet, 27 | Writer.write_octet, 1 << 12), 28 | ('priority', 'octet', Reader.read_octet, Writer.write_octet, 1 << 11), 29 | ('correlation_id', 'shortstr', Reader.read_shortstr, 30 | Writer.write_shortstr, 1 << 10), 31 | ('reply_to', 'shortstr', Reader.read_shortstr, 32 | Writer.write_shortstr, 1 << 9), 33 | ('expiration', 'shortstr', Reader.read_shortstr, 34 | Writer.write_shortstr, 1 << 8), 35 | ('message_id', 'shortstr', Reader.read_shortstr, 36 | Writer.write_shortstr, 1 << 7), 37 | ('timestamp', 'timestamp', Reader.read_timestamp, 38 | Writer.write_timestamp, 1 << 6), 39 | ('type', 'shortstr', Reader.read_shortstr, 40 | Writer.write_shortstr, 1 << 5), 41 | ('user_id', 'shortstr', Reader.read_shortstr, 42 | Writer.write_shortstr, 1 << 4), 43 | ('app_id', 'shortstr', Reader.read_shortstr, 44 | Writer.write_shortstr, 1 << 3), 45 | ('cluster_id', 'shortstr', Reader.read_shortstr, 46 | Writer.write_shortstr, 1 << 2) 47 | ] 48 | DEFAULT_PROPERTIES = True 49 | 50 | @classmethod 51 | def type(cls): 52 | return 2 53 | 54 | @property 55 | def class_id(self): 56 | return self._class_id 57 | 58 | @property 59 | def weight(self): 60 | return self._weight 61 | 62 | @property 63 | def size(self): 64 | return self._size 65 | 66 | @property 67 | def properties(self): 68 | return self._properties 69 | 70 | @classmethod 71 | def parse(self, channel_id, payload): 72 | ''' 73 | Parse a header frame for a channel given a Reader payload. 74 | ''' 75 | class_id = payload.read_short() 76 | weight = payload.read_short() 77 | size = payload.read_longlong() 78 | properties = {} 79 | 80 | # The AMQP spec is overly-complex when it comes to handling header 81 | # frames. The spec says that in addition to the first 16bit field, 82 | # additional ones can follow which /may/ then be in the property list 83 | # (because bit flags aren't in the list). Properly implementing custom 84 | # values requires the ability change the properties and their types, 85 | # which someone is welcome to do, but seriously, what's the point? 86 | # Because the complexity of parsing and writing this frame directly 87 | # impacts the speed at which messages can be processed, there are two 88 | # branches for both a fast parse which assumes no changes to the 89 | # properties and a slow parse. For now it's up to someone using custom 90 | # headers to flip the flag. 91 | if self.DEFAULT_PROPERTIES: 92 | flag_bits = payload.read_short() 93 | for key, proptype, rfunc, wfunc, mask in self.PROPERTIES: 94 | if flag_bits & mask: 95 | properties[key] = rfunc(payload) 96 | else: 97 | flags = [] 98 | while True: 99 | flag_bits = payload.read_short() 100 | flags.append(flag_bits) 101 | if flag_bits & 1 == 0: 102 | break 103 | 104 | shift = 0 105 | for key, proptype, rfunc, wfunc, mask in self.PROPERTIES: 106 | if shift == 0: 107 | if not flags: 108 | break 109 | flag_bits, flags = flags[0], flags[1:] 110 | shift = 15 111 | if flag_bits & (1 << shift): 112 | properties[key] = rfunc(payload) 113 | shift -= 1 114 | 115 | return HeaderFrame(channel_id, class_id, weight, size, properties) 116 | 117 | def __init__(self, channel_id, class_id, weight, size, properties={}): 118 | Frame.__init__(self, channel_id) 119 | self._class_id = class_id 120 | self._weight = weight 121 | self._size = size 122 | self._properties = properties 123 | 124 | def __str__(self): 125 | return "%s[channel: %d, class_id: %d, weight: %d, size: %d, properties: %s]" % ( 126 | self.__class__.__name__, self.channel_id, self._class_id, 127 | self._weight, self._size, self._properties) 128 | 129 | def write_frame(self, buf): 130 | ''' 131 | Write the frame into an existing buffer. 132 | ''' 133 | writer = Writer(buf) 134 | writer.write_octet(self.type()) 135 | writer.write_short(self.channel_id) 136 | 137 | # Track the position where we're going to write the total length 138 | # of the frame arguments. 139 | stream_args_len_pos = len(buf) 140 | writer.write_long(0) 141 | 142 | stream_method_pos = len(buf) 143 | 144 | writer.write_short(self._class_id) 145 | writer.write_short(self._weight) 146 | writer.write_longlong(self._size) 147 | 148 | # Like frame parsing, branch to faster code for default properties 149 | if self.DEFAULT_PROPERTIES: 150 | # Track the position where we're going to write the flags. 151 | flags_pos = len(buf) 152 | writer.write_short(0) 153 | flag_bits = 0 154 | for key, proptype, rfunc, wfunc, mask in self.PROPERTIES: 155 | val = self._properties.get(key, None) 156 | if val is not None: 157 | flag_bits |= mask 158 | wfunc(writer, val) 159 | writer.write_short_at(flag_bits, flags_pos) 160 | else: 161 | shift = 15 162 | flag_bits = 0 163 | flags = [] 164 | stack = deque() 165 | for key, proptype, rfunc, wfunc, mask in self.PROPERTIES: 166 | val = self._properties.get(key, None) 167 | if val is not None: 168 | if shift == 0: 169 | flags.append(flag_bits) 170 | flag_bits = 0 171 | shift = 15 172 | 173 | flag_bits |= (1 << shift) 174 | stack.append((wfunc, val)) 175 | 176 | shift -= 1 177 | 178 | flags.append(flag_bits) 179 | for flag_bits in flags: 180 | writer.write_short(flag_bits) 181 | for method, val in stack: 182 | method(writer, val) 183 | 184 | # Write the total length back at the beginning of the frame 185 | stream_len = len(buf) - stream_method_pos 186 | writer.write_long_at(stream_len, stream_args_len_pos) 187 | 188 | writer.write_octet(0xce) 189 | 190 | HeaderFrame.register() 191 | -------------------------------------------------------------------------------- /haigha/frames/heartbeat_frame.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.frames.frame import Frame 8 | from haigha.writer import Writer 9 | 10 | 11 | class HeartbeatFrame(Frame): 12 | 13 | ''' 14 | Frame for heartbeats. 15 | ''' 16 | 17 | @classmethod 18 | def type(cls): 19 | # NOTE: The PDF spec say this should be 4 but the xml spec say it 20 | # should be 8 RabbitMQ seems to implement this as 8, but maybe that's 21 | # a difference between 0.8 and 0.9 protocols 22 | # Using Rabbit 2.1.1 and protocol 0.9.1 it seems that 8 is indeed 23 | # the correct type @AW 24 | # PDF spec: http://www.amqp.org/confluence/download/attachments/720900/amqp0-9-1.pdf?version=1&modificationDate=1227526523000 25 | # XML spec: http://www.amqp.org/confluence/download/attachments/720900/amqp0-9-1.xml?version=1&modificationDate=1227526672000 26 | # This is addressed in 27 | # http://dev.rabbitmq.com/wiki/Amqp091Errata#section_29 28 | return 8 29 | 30 | @classmethod 31 | def parse(self, channel_id, payload): 32 | return HeartbeatFrame(channel_id) 33 | 34 | def write_frame(self, buf): 35 | writer = Writer(buf) 36 | writer.write_octet(self.type()) 37 | writer.write_short(self.channel_id) 38 | writer.write_long(0) 39 | writer.write_octet(0xce) 40 | 41 | HeartbeatFrame.register() 42 | -------------------------------------------------------------------------------- /haigha/frames/method_frame.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.frames.frame import Frame 8 | from haigha.reader import Reader 9 | from haigha.writer import Writer 10 | 11 | 12 | class MethodFrame(Frame): 13 | 14 | ''' 15 | Frame which carries identifier for methods. 16 | ''' 17 | 18 | @classmethod 19 | def type(cls): 20 | return 1 21 | 22 | @property 23 | def class_id(self): 24 | return self._class_id 25 | 26 | @property 27 | def method_id(self): 28 | return self._method_id 29 | 30 | @property 31 | def args(self): 32 | return self._args 33 | 34 | @classmethod 35 | def parse(self, channel_id, payload): 36 | class_id = payload.read_short() 37 | method_id = payload.read_short() 38 | return MethodFrame(channel_id, class_id, method_id, payload) 39 | 40 | def __init__(self, channel_id, class_id, method_id, args=None): 41 | Frame.__init__(self, channel_id) 42 | self._class_id = class_id 43 | self._method_id = method_id 44 | self._args = args 45 | 46 | def __str__(self): 47 | if isinstance(self.args, (Reader, Writer)): 48 | return "%s[channel: %d, class_id: %d, method_id: %d, args: %s]" %\ 49 | (self.__class__.__name__, self.channel_id, 50 | self.class_id, self.method_id, str(self.args)) 51 | else: 52 | return "%s[channel: %d, class_id: %d, method_id: %d, args: None]" %\ 53 | (self.__class__.__name__, self.channel_id, 54 | self.class_id, self.method_id) 55 | 56 | def write_frame(self, buf): 57 | writer = Writer(buf) 58 | writer.write_octet(self.type()) 59 | writer.write_short(self.channel_id) 60 | 61 | # Write a temporary value for the total length of the frame 62 | stream_args_len_pos = len(buf) 63 | writer.write_long(0) 64 | 65 | # Mark the point in the stream where we start writing arguments, 66 | # *including* the class and method ids. 67 | stream_method_pos = len(buf) 68 | 69 | writer.write_short(self.class_id) 70 | writer.write_short(self.method_id) 71 | 72 | # This is assuming that args is a Writer 73 | if self._args is not None: 74 | writer.write(self._args.buffer()) 75 | 76 | # Write the total length back at the position we allocated 77 | stream_len = len(buf) - stream_method_pos 78 | writer.write_long_at(stream_len, stream_args_len_pos) 79 | 80 | # Write the footer 81 | writer.write_octet(0xce) 82 | 83 | 84 | MethodFrame.register() 85 | -------------------------------------------------------------------------------- /haigha/message.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | 8 | class Message(object): 9 | 10 | ''' 11 | Represents an AMQP message. 12 | ''' 13 | 14 | def __init__(self, body='', delivery_info=None, return_info=None, 15 | **properties): 16 | ''' 17 | :param delivery_info: pass only if messages was received via 18 | basic.deliver or basic.get_ok; MUST be None otherwise; default: None 19 | :param return_info: pass only if message was returned via basic.return; 20 | MUST be None otherwise; default: None 21 | ''' 22 | if isinstance(body, unicode): 23 | if 'content_encoding' not in properties: 24 | properties['content_encoding'] = 'utf-8' 25 | body = body.encode(properties['content_encoding']) 26 | 27 | if not isinstance(body, (str, unicode, bytearray)): 28 | raise TypeError("Invalid message content type %s" % (type(body))) 29 | 30 | self._body = body 31 | self._delivery_info = delivery_info 32 | self._return_info = return_info 33 | self._properties = properties 34 | 35 | @property 36 | def body(self): 37 | return self._body 38 | 39 | def __len__(self): 40 | return len(self._body) 41 | 42 | def __nonzero__(self): 43 | '''Have to define this because length is defined.''' 44 | return True 45 | 46 | def __eq__(self, other): 47 | if isinstance(other, Message): 48 | return self._properties == other._properties and \ 49 | self._body == other._body 50 | return False 51 | 52 | @property 53 | def delivery_info(self): 54 | '''delivery_info dict if message was received via basic.deliver or 55 | basic.get_ok; None otherwise. 56 | ''' 57 | return self._delivery_info 58 | 59 | @property 60 | def return_info(self): 61 | '''return_info dict if message was returned via basic.return; None 62 | otherwise. 63 | 64 | properties: 65 | 'channel': Channel instance 66 | 'reply_code': reply code (int) 67 | 'reply_text': reply text 68 | 'exchange': exchange name 69 | 'routing_key': routing key 70 | ''' 71 | return self._return_info 72 | 73 | @property 74 | def properties(self): 75 | return self._properties 76 | 77 | def __str__(self): 78 | return ("Message[body: %s, delivery_info: %s, return_info: %s, " 79 | "properties: %s]") %\ 80 | (str(self._body).encode('string_escape'), 81 | self._delivery_info, self.return_info, self._properties) 82 | -------------------------------------------------------------------------------- /haigha/transports/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/haigha/transports/__init__.py -------------------------------------------------------------------------------- /haigha/transports/event_transport.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | import warnings 8 | 9 | from haigha.transports.transport import Transport 10 | 11 | try: 12 | from eventsocket import EventSocket 13 | import event 14 | except ImportError: 15 | warnings.warn('Failed to load EventSocket and event modules') 16 | EventSocket = None 17 | event = None 18 | 19 | 20 | class EventTransport(Transport): 21 | 22 | ''' 23 | Transport using libevent-based EventSocket. 24 | ''' 25 | 26 | def __init__(self, *args): 27 | super(EventTransport, self).__init__(*args) 28 | self._synchronous = False 29 | 30 | ### 31 | # EventSocket callbacks 32 | ### 33 | def _sock_close_cb(self, sock): 34 | self._connection.transport_closed( 35 | msg='socket to %s closed unexpectedly' % (self._host), 36 | ) 37 | 38 | def _sock_error_cb(self, sock, msg, exception=None): 39 | self._connection.transport_closed( 40 | msg='error on connection to %s: %s' % (self._host, msg) 41 | ) 42 | 43 | def _sock_read_cb(self, sock): 44 | self.connection.read_frames() 45 | 46 | ### 47 | # Transport API 48 | ### 49 | def connect(self, (host, port)): 50 | ''' 51 | Connect assuming a host and port tuple. Implemented as non-blocking, 52 | and will close the transport if there's an error 53 | ''' 54 | self._host = "%s:%s" % (host, port) 55 | self._sock = EventSocket( 56 | read_cb=self._sock_read_cb, 57 | close_cb=self._sock_close_cb, 58 | error_cb=self._sock_error_cb, 59 | debug=self.connection.debug, 60 | logger=self.connection.logger) 61 | if self.connection._sock_opts: 62 | for k, v in self.connection._sock_opts.iteritems(): 63 | family, type = k 64 | self._sock.setsockopt(family, type, v) 65 | self._sock.setblocking(False) 66 | self._sock.connect( 67 | (host, port), timeout=self.connection._connect_timeout) 68 | self._heartbeat_timeout = None 69 | 70 | def read(self, timeout=None): 71 | ''' 72 | Read from the transport. If no data is available, should return None. 73 | The timeout is ignored as this returns only data that has already 74 | been buffered locally. 75 | ''' 76 | # NOTE: copying over this comment from Connection, because there is 77 | # knowledge captured here, even if the details are stale 78 | # Because of the timer callback to dataRead when we re-buffered, 79 | # there's a chance that in between we've lost the socket. If that's 80 | # the case, just silently return as some code elsewhere would have 81 | # already notified us. That bug could be fixed by improving the 82 | # message reading so that we consume all possible messages and ensure 83 | # that only a partial message was rebuffered, so that we can rely on 84 | # the next read event to read the subsequent message. 85 | if not hasattr(self, '_sock'): 86 | return None 87 | 88 | # This is sort of a hack because we're faking that data is ready, but 89 | # it works for purposes of supporting timeouts 90 | if timeout: 91 | if self._heartbeat_timeout: 92 | self._heartbeat_timeout.delete() 93 | self._heartbeat_timeout = \ 94 | event.timeout(timeout, self._sock_read_cb, self._sock) 95 | elif self._heartbeat_timeout: 96 | self._heartbeat_timeout.delete() 97 | self._heartbeat_timeout = None 98 | 99 | return self._sock.read() 100 | 101 | def buffer(self, data): 102 | ''' 103 | Buffer unused bytes from the input stream. 104 | ''' 105 | if not hasattr(self, '_sock'): 106 | return None 107 | self._sock.buffer(data) 108 | 109 | def write(self, data): 110 | ''' 111 | Write some bytes to the transport. 112 | ''' 113 | if not hasattr(self, '_sock'): 114 | return 115 | self._sock.write(data) 116 | 117 | def disconnect(self): 118 | ''' 119 | Disconnect from the transport. Typically socket.close(). This call is 120 | welcome to raise exceptions, which the Connection will catch. 121 | 122 | The transport is encouraged to allow for any pending writes to complete 123 | before closing the socket. 124 | ''' 125 | if not hasattr(self, '_sock'): 126 | return 127 | 128 | # TODO: If there are bytes left on the output, queue the close for 129 | # later. 130 | self._sock.close_cb = None 131 | self._sock.close() 132 | -------------------------------------------------------------------------------- /haigha/transports/gevent_transport.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | import warnings 8 | 9 | from haigha.transports.socket_transport import SocketTransport 10 | 11 | try: 12 | import gevent 13 | from gevent.event import Event 14 | try: 15 | # Semaphore moved here since gevent-1.0b2 16 | from gevent.lock import Semaphore 17 | except ImportError: 18 | from gevent.coros import Semaphore 19 | from gevent import socket 20 | from gevent import pool 21 | except ImportError: 22 | warnings.warn('Failed to load gevent modules') 23 | gevent = None 24 | Event = None 25 | Semaphore = None 26 | socket = None 27 | pool = None 28 | 29 | class GeventTransport(SocketTransport): 30 | 31 | ''' 32 | Transport using gevent backend. It relies on gevent's implementation of 33 | sendall to send whole frames at a time. On the input side, it uses a gevent 34 | semaphore to ensure exclusive access to the socket and input buffer. 35 | ''' 36 | 37 | def __init__(self, *args, **kwargs): 38 | super(GeventTransport, self).__init__(*args) 39 | 40 | self._synchronous = False 41 | self._read_lock = Semaphore() 42 | self._write_lock = Semaphore() 43 | self._read_wait = Event() 44 | 45 | ### 46 | # Transport API 47 | ### 48 | 49 | def connect(self, (host, port)): 50 | ''' 51 | Connect using a host,port tuple 52 | ''' 53 | super(GeventTransport, self).connect((host, port), klass=socket.socket) 54 | 55 | def read(self, timeout=None): 56 | ''' 57 | Read from the transport. If no data is available, should return None. 58 | If timeout>0, will only block for `timeout` seconds. 59 | ''' 60 | # If currently locked, another greenlet is trying to read, so yield 61 | # control and then return none. Required if a Connection is configured 62 | # to be synchronous, a sync callback is trying to read, and there's 63 | # another read loop running read_frames. Without it, the run loop will 64 | # release the lock but then immediately acquire it again. Yielding 65 | # control in the reading thread after bytes are read won't fix 66 | # anything, because it's quite possible the bytes read resulted in a 67 | # frame that satisfied the synchronous callback, and so this needs to 68 | # return immediately to first check the current status of synchronous 69 | # callbacks before attempting to read again. 70 | if self._read_lock.locked(): 71 | self._read_wait.wait(timeout) 72 | return None 73 | 74 | self._read_lock.acquire() 75 | try: 76 | return super(GeventTransport, self).read(timeout=timeout) 77 | finally: 78 | self._read_lock.release() 79 | self._read_wait.set() 80 | self._read_wait.clear() 81 | 82 | def buffer(self, data): 83 | ''' 84 | Buffer unused bytes from the input stream. 85 | ''' 86 | self._read_lock.acquire() 87 | try: 88 | return super(GeventTransport, self).buffer(data) 89 | finally: 90 | self._read_lock.release() 91 | 92 | def write(self, data): 93 | ''' 94 | Write some bytes to the transport. 95 | ''' 96 | # MUST use a lock here else gevent could raise an exception if 2 97 | # greenlets try to write at the same time. I was hoping that 98 | # sendall() would do that blocking for me, but I guess not. May 99 | # require an eventsocket-like buffer to speed up under high load. 100 | self._write_lock.acquire() 101 | try: 102 | return super(GeventTransport, self).write(data) 103 | finally: 104 | self._write_lock.release() 105 | 106 | 107 | class GeventPoolTransport(GeventTransport): 108 | 109 | def __init__(self, *args, **kwargs): 110 | super(GeventPoolTransport, self).__init__(*args) 111 | 112 | self._pool = kwargs.get('pool', None) 113 | if not self._pool: 114 | self._pool = gevent.pool.Pool() 115 | 116 | @property 117 | def pool(self): 118 | '''Get a handle to the gevent pool.''' 119 | return self._pool 120 | 121 | def process_channels(self, channels): 122 | ''' 123 | Process a set of channels by calling Channel.process_frames() on each. 124 | Some transports may choose to do this in unique ways, such as through 125 | a pool of threads. 126 | 127 | The default implementation will simply iterate over them and call 128 | process_frames() on each. 129 | ''' 130 | for channel in channels: 131 | self._pool.spawn(channel.process_frames) 132 | -------------------------------------------------------------------------------- /haigha/transports/socket_transport.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from haigha.transports.transport import Transport 8 | 9 | import errno 10 | import socket 11 | 12 | 13 | class SocketTransport(Transport): 14 | 15 | ''' 16 | A simple blocking socket transport. 17 | ''' 18 | 19 | def __init__(self, *args): 20 | super(SocketTransport, self).__init__(*args) 21 | self._synchronous = True 22 | self._buffer = bytearray() 23 | 24 | ### 25 | # Transport API 26 | ### 27 | def connect(self, (host, port), klass=socket.socket): 28 | '''Connect assuming a host and port tuple. 29 | 30 | :param tuple: A tuple containing host and port for a connection. 31 | :param klass: A implementation of socket.socket. 32 | :raises socket.gaierror: If no address can be resolved. 33 | :raises socket.error: If no connection can be made. 34 | ''' 35 | self._host = "%s:%s" % (host, port) 36 | 37 | for info in socket.getaddrinfo(host, port, 0, 0, socket.IPPROTO_TCP): 38 | 39 | family, socktype, proto, _, sockaddr = info 40 | self._sock = klass(family, socktype, proto) 41 | self._sock.settimeout(self.connection._connect_timeout) 42 | if self.connection._sock_opts: 43 | _sock_opts = self.connection._sock_opts 44 | for (level, optname), value in _sock_opts.iteritems(): 45 | self._sock.setsockopt(level, optname, value) 46 | try: 47 | 48 | self._sock.connect(sockaddr) 49 | 50 | except socket.error: 51 | 52 | self.connection.logger.exception( 53 | "Failed to connect to %s:", 54 | sockaddr, 55 | ) 56 | continue 57 | 58 | # After connecting, switch to full-blocking mode. 59 | self._sock.settimeout(None) 60 | break 61 | 62 | else: 63 | 64 | raise 65 | 66 | def read(self, timeout=None): 67 | ''' 68 | Read from the transport. If timeout>0, will only block for `timeout` 69 | seconds. 70 | ''' 71 | e = None 72 | if not hasattr(self, '_sock'): 73 | return None 74 | 75 | try: 76 | # Note that we ignore both None and 0, i.e. we either block with a 77 | # timeout or block completely and let gevent sort it out. 78 | if timeout: 79 | self._sock.settimeout(timeout) 80 | else: 81 | self._sock.settimeout(None) 82 | data = self._sock.recv( 83 | self._sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF)) 84 | 85 | if len(data): 86 | if self.connection.debug > 1: 87 | self.connection.logger.debug( 88 | 'read %d bytes from %s' % (len(data), self._host)) 89 | if len(self._buffer): 90 | self._buffer.extend(data) 91 | data = self._buffer 92 | self._buffer = bytearray() 93 | return data 94 | 95 | # Note that no data means the socket is closed and we'll mark that 96 | # below 97 | 98 | except socket.timeout as e: 99 | # Note that this is implemented differently and though it would be 100 | # caught as an EnvironmentError, it has no errno. Not sure whose 101 | # fault that is. 102 | return None 103 | 104 | except EnvironmentError as e: 105 | # thrown if we have a timeout and no data 106 | if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK, errno.EINTR): 107 | return None 108 | 109 | self.connection.logger.exception( 110 | 'error reading from %s' % (self._host)) 111 | 112 | self.connection.transport_closed( 113 | msg='error reading from %s' % (self._host)) 114 | if e: 115 | raise 116 | 117 | def buffer(self, data): 118 | ''' 119 | Buffer unused bytes from the input stream. 120 | ''' 121 | if not hasattr(self, '_sock'): 122 | return None 123 | 124 | # data will always be a byte array 125 | if len(self._buffer): 126 | self._buffer.extend(data) 127 | else: 128 | self._buffer = bytearray(data) 129 | 130 | def write(self, data): 131 | ''' 132 | Write some bytes to the transport. 133 | ''' 134 | if not hasattr(self, '_sock'): 135 | return None 136 | 137 | try: 138 | self._sock.sendall(data) 139 | 140 | if self.connection.debug > 1: 141 | self.connection.logger.debug( 142 | 'sent %d bytes to %s' % (len(data), self._host)) 143 | 144 | return 145 | except EnvironmentError: 146 | # sockets raise this type of error, and since if sendall() fails 147 | # we're left in an indeterminate state, assume that any error we 148 | # catch means that the connection is dead. Note that this 149 | # assumption requires this to be a blocking socket; if we ever 150 | # support non-blocking in this class then this whole method has 151 | # to change a lot. 152 | self.connection.logger.exception( 153 | 'error writing to %s' % (self._host)) 154 | 155 | self.connection.transport_closed( 156 | msg='error writing to %s' % (self._host)) 157 | 158 | def disconnect(self): 159 | ''' 160 | Disconnect from the transport. Typically socket.close(). This call is 161 | welcome to raise exceptions, which the Connection will catch. 162 | ''' 163 | if not hasattr(self, '_sock'): 164 | return None 165 | 166 | try: 167 | self._sock.close() 168 | finally: 169 | self._sock = None 170 | -------------------------------------------------------------------------------- /haigha/transports/transport.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | 8 | class Transport(object): 9 | 10 | ''' 11 | Base class and API for Transports 12 | ''' 13 | 14 | def __init__(self, connection): 15 | ''' 16 | Initialize a transport on a haigha.Connection instance. 17 | ''' 18 | self._connection = connection 19 | 20 | @property 21 | def synchronous(self): 22 | '''Return True if this is a synchronous transport, False otherwise.''' 23 | # Note that subclasses must define this. 24 | return self._synchronous 25 | 26 | @property 27 | def connection(self): 28 | return self._connection 29 | 30 | def process_channels(self, channels): 31 | ''' 32 | Process a set of channels by calling Channel.process_frames() on each. 33 | Some transports may choose to do this in unique ways, such as through 34 | a pool of threads. 35 | 36 | The default implementation will simply iterate over them and call 37 | process_frames() on each. 38 | ''' 39 | for channel in channels: 40 | channel.process_frames() 41 | 42 | def read(self, timeout=None): 43 | ''' 44 | Read from the transport. If no data is available, should return None. 45 | The return value can be any data type that is supported by the 46 | haigha.Reader class. 47 | 48 | Caller passes in an optional timeout. Each transport determines how to 49 | implement this. 50 | ''' 51 | return None 52 | 53 | def buffer(self, data): 54 | ''' 55 | Buffer unused bytes from the input stream. 56 | ''' 57 | 58 | def write(self, data): 59 | ''' 60 | Write some bytes to the transport. 61 | ''' 62 | 63 | def disconnect(self): 64 | ''' 65 | Disconnect from the transport. Typically socket.close(). This call is 66 | welcome to raise exceptions, which the Connection will catch. 67 | 68 | The transport is encouraged to allow for any pending writes to complete 69 | before closing the socket. 70 | ''' 71 | -------------------------------------------------------------------------------- /haigha/writer.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from struct import Struct 8 | from calendar import timegm 9 | from datetime import datetime 10 | from decimal import Decimal 11 | from operator import xor 12 | 13 | 14 | class Writer(object): 15 | 16 | """ 17 | Implements writing of structured AMQP data. Buffers data directly to a 18 | bytearray or a buffer supplied in the constructor. The buffer must 19 | supply append, extend and struct.pack_into semantics. 20 | """ 21 | 22 | def __init__(self, buf=None): 23 | if buf is not None: 24 | self._output_buffer = buf 25 | else: 26 | self._output_buffer = bytearray() 27 | 28 | def __str__(self): 29 | return ''.join([ 30 | '\\x%s' % (chr(c).encode('hex')) for c in self._output_buffer]) 31 | 32 | __repr__ = __str__ 33 | 34 | def __eq__(self, other): 35 | if isinstance(other, Writer): 36 | return self._output_buffer == other._output_buffer 37 | return False 38 | 39 | def buffer(self): 40 | ''' 41 | Get the buffer that this has written to. Returns bytearray. 42 | ''' 43 | return self._output_buffer 44 | 45 | def write(self, s): 46 | """ 47 | Write a plain Python string, with no special encoding. 48 | """ 49 | self._output_buffer.extend(s) 50 | return self 51 | 52 | def write_bits(self, *args): 53 | ''' 54 | Write multiple bits in a single byte field. The bits will be written in 55 | little-endian order, but should be supplied in big endian order. Will 56 | raise ValueError when more than 8 arguments are supplied. 57 | 58 | write_bits(True, False) => 0x02 59 | ''' 60 | # Would be nice to make this a bit smarter 61 | if len(args) > 8: 62 | raise ValueError("Can only write 8 bits at a time") 63 | 64 | self._output_buffer.append(chr( 65 | reduce(lambda x, y: xor(x, args[y] << y), xrange(len(args)), 0))) 66 | 67 | return self 68 | 69 | def write_bit(self, b, pack=Struct('B').pack): 70 | ''' 71 | Write a single bit. Convenience method for single bit args. 72 | ''' 73 | self._output_buffer.append(pack(True if b else False)) 74 | return self 75 | 76 | def write_octet(self, n, pack=Struct('B').pack): 77 | """ 78 | Write an integer as an unsigned 8-bit value. 79 | """ 80 | if 0 <= n <= 255: 81 | self._output_buffer.append(pack(n)) 82 | else: 83 | raise ValueError('Octet %d out of range 0..255', n) 84 | return self 85 | 86 | def write_short(self, n, pack=Struct('>H').pack): 87 | """ 88 | Write an integer as an unsigned 16-bit value. 89 | """ 90 | if 0 <= n <= 0xFFFF: 91 | self._output_buffer.extend(pack(n)) 92 | else: 93 | raise ValueError('Short %d out of range 0..0xFFFF', n) 94 | return self 95 | 96 | def write_short_at(self, n, pos, pack_into=Struct('>H').pack_into): 97 | ''' 98 | Write an unsigned 16bit value at a specific position in the buffer. 99 | Used for writing tables and frames. 100 | ''' 101 | if 0 <= n <= 0xFFFF: 102 | pack_into(self._output_buffer, pos, n) 103 | else: 104 | raise ValueError('Short %d out of range 0..0xFFFF', n) 105 | return self 106 | 107 | def write_long(self, n, pack=Struct('>I').pack): 108 | """ 109 | Write an integer as an unsigned 32-bit value. 110 | """ 111 | if 0 <= n <= 0xFFFFFFFF: 112 | self._output_buffer.extend(pack(n)) 113 | else: 114 | raise ValueError('Long %d out of range 0..0xFFFFFFFF', n) 115 | return self 116 | 117 | def write_long_at(self, n, pos, pack_into=Struct('>I').pack_into): 118 | ''' 119 | Write an unsigned 32bit value at a specific position in the buffer. 120 | Used for writing tables and frames. 121 | ''' 122 | if 0 <= n <= 0xFFFFFFFF: 123 | pack_into(self._output_buffer, pos, n) 124 | else: 125 | raise ValueError('Long %d out of range 0..0xFFFFFFFF', n) 126 | return self 127 | 128 | def write_longlong(self, n, pack=Struct('>Q').pack): 129 | """ 130 | Write an integer as an unsigned 64-bit value. 131 | """ 132 | if 0 <= n <= 0xFFFFFFFFFFFFFFFF: 133 | self._output_buffer.extend(pack(n)) 134 | else: 135 | raise ValueError( 136 | 'Longlong %d out of range 0..0xFFFFFFFFFFFFFFFF', n) 137 | return self 138 | 139 | def write_shortstr(self, s): 140 | """ 141 | Write a string up to 255 bytes long after encoding. If passed 142 | a unicode string, encode as UTF-8. 143 | """ 144 | if isinstance(s, unicode): 145 | s = s.encode('utf-8') 146 | self.write_octet(len(s)) 147 | self.write(s) 148 | return self 149 | 150 | def write_longstr(self, s): 151 | """ 152 | Write a string up to 2**32 bytes long after encoding. If passed 153 | a unicode string, encode as UTF-8. 154 | """ 155 | if isinstance(s, unicode): 156 | s = s.encode('utf-8') 157 | self.write_long(len(s)) 158 | self.write(s) 159 | return self 160 | 161 | def write_timestamp(self, t, pack=Struct('>Q').pack): 162 | """ 163 | Write out a Python datetime.datetime object as a 64-bit integer 164 | representing seconds since the Unix UTC epoch. 165 | """ 166 | # Double check timestamp, can't imagine why it would be signed 167 | self._output_buffer.extend(pack(long(timegm(t.timetuple())))) 168 | return self 169 | 170 | # NOTE: coding to http://dev.rabbitmq.com/wiki/Amqp091Errata#section_3 and 171 | # NOT spec 0.9.1. It seems that Rabbit and other brokers disagree on this 172 | # section for now. 173 | def write_table(self, d): 174 | """ 175 | Write out a Python dictionary made of up string keys, and values 176 | that are strings, signed integers, Decimal, datetime.datetime, or 177 | sub-dictionaries following the same constraints. 178 | """ 179 | # HACK: encoding of AMQP tables is broken because it requires the 180 | # length of the /encoded/ data instead of the number of items. To 181 | # support streaming, fiddle with cursor position, rewinding to write 182 | # the real length of the data. Generally speaking, I'm not a fan of 183 | # the AMQP encoding scheme, it could be much faster. 184 | table_len_pos = len(self._output_buffer) 185 | self.write_long(0) 186 | table_data_pos = len(self._output_buffer) 187 | 188 | for key, value in d.iteritems(): 189 | self._write_item(key, value) 190 | 191 | table_end_pos = len(self._output_buffer) 192 | table_len = table_end_pos - table_data_pos 193 | 194 | self.write_long_at(table_len, table_len_pos) 195 | return self 196 | 197 | def _write_item(self, key, value): 198 | self.write_shortstr(key) 199 | self._write_field(value) 200 | 201 | def _write_field(self, value): 202 | writer = self.field_type_map.get(type(value)) 203 | if writer: 204 | writer(self, value) 205 | else: 206 | for kls, writer in self.field_type_map.items(): 207 | if isinstance(value, kls): 208 | writer(self, value) 209 | break 210 | else: 211 | # Write a None because we've already written a key 212 | self._field_none(value) 213 | 214 | def _field_bool(self, val, pack=Struct('B').pack): 215 | self._output_buffer.append('t') 216 | self._output_buffer.append(pack(True if val else False)) 217 | 218 | def _field_int(self, val, short_pack=Struct('>h').pack, 219 | int_pack=Struct('>i').pack, long_pack=Struct('>q').pack): 220 | if -2 ** 15 <= val < 2 ** 15: 221 | self._output_buffer.append('s') 222 | self._output_buffer.extend(short_pack(val)) 223 | elif -2 ** 31 <= val < 2 ** 31: 224 | self._output_buffer.append('I') 225 | self._output_buffer.extend(int_pack(val)) 226 | else: 227 | self._output_buffer.append('l') 228 | self._output_buffer.extend(long_pack(val)) 229 | 230 | def _field_double(self, val, pack=Struct('>d').pack): 231 | self._output_buffer.append('d') 232 | self._output_buffer.extend(pack(val)) 233 | 234 | # Coding to http://dev.rabbitmq.com/wiki/Amqp091Errata#section_3 which 235 | # differs from spec in that the value is signed. 236 | def _field_decimal(self, val, exp_pack=Struct('B').pack, 237 | dig_pack=Struct('>i').pack): 238 | self._output_buffer.append('D') 239 | sign, digits, exponent = val.as_tuple() 240 | v = 0 241 | for d in digits: 242 | v = (v * 10) + d 243 | if sign: 244 | v = -v 245 | self._output_buffer.append(exp_pack(-exponent)) 246 | self._output_buffer.extend(dig_pack(v)) 247 | 248 | def _field_str(self, val): 249 | self._output_buffer.append('S') 250 | self.write_longstr(val) 251 | 252 | def _field_unicode(self, val): 253 | val = val.encode('utf-8') 254 | self._output_buffer.append('S') 255 | self.write_longstr(val) 256 | 257 | def _field_timestamp(self, val): 258 | self._output_buffer.append('T') 259 | self.write_timestamp(val) 260 | 261 | def _field_table(self, val): 262 | self._output_buffer.append('F') 263 | self.write_table(val) 264 | 265 | def _field_none(self, val): 266 | self._output_buffer.append('V') 267 | 268 | def _field_bytearray(self, val): 269 | self._output_buffer.append('x') 270 | self.write_longstr(val) 271 | 272 | def _field_iterable(self, val): 273 | self._output_buffer.append('A') 274 | for x in val: 275 | self._write_field(x) 276 | 277 | field_type_map = { 278 | bool: _field_bool, 279 | int: _field_int, 280 | long: _field_int, 281 | float: _field_double, 282 | Decimal: _field_decimal, 283 | str: _field_str, 284 | unicode: _field_unicode, 285 | datetime: _field_timestamp, 286 | dict: _field_table, 287 | type(None): _field_none, 288 | bytearray: _field_bytearray, 289 | } 290 | 291 | # 0.9.1 spec mapping 292 | # field_type_map = { 293 | # bool : _field_bool, 294 | # int : _field_int, 295 | # long : _field_int, 296 | # float : _field_double, 297 | # Decimal : _field_decimal, 298 | # str : _field_str, 299 | # unicode : _field_unicode, 300 | # datetime : _field_timestamp, 301 | # dict : _field_table, 302 | # None : _field_none, 303 | # bytearray : _field_bytearray, 304 | # list : _field_iterable, 305 | # tuple : _field_iterable, 306 | # set : _field_iterable, 307 | # } 308 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/requirements.txt -------------------------------------------------------------------------------- /scripts/console: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #-*- coding:utf-8 -*- 3 | 4 | import sys, os 5 | sys.path.append(os.path.abspath(".")) 6 | sys.path.append(os.path.abspath("..")) 7 | 8 | import logging 9 | import random 10 | import socket 11 | from optparse import OptionParser 12 | 13 | from haigha.connection import Connection 14 | from haigha.message import Message 15 | 16 | parser = OptionParser( 17 | usage='Usage: synchronous_test [options]' 18 | ) 19 | parser.add_option('--user', default='guest', type='string') 20 | parser.add_option('--pass', default='guest', dest='password', type='string') 21 | parser.add_option('--vhost', default='/', type='string') 22 | parser.add_option('--host', default='localhost', type='string') 23 | parser.add_option('--debug', default=0, action='count') 24 | 25 | (options,args) = parser.parse_args() 26 | 27 | debug = options.debug 28 | level = logging.DEBUG if debug else logging.INFO 29 | 30 | # Setup logging 31 | logging.basicConfig(level=level, format="%(message)s", stream=sys.stdout ) 32 | 33 | sock_opts = { 34 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 35 | } 36 | 37 | connection = Connection(logger=logging, debug=debug, 38 | user=options.user, password=options.password, 39 | vhost=options.vhost, host=options.host, 40 | heartbeat=None, 41 | sock_opts=sock_opts, 42 | transport='socket') 43 | 44 | hello = '''Welcome to haigha console. 45 | 46 | `connection` - The established connection to %s:%s 47 | `Message` - The Message type 48 | 49 | Usage: Create a channel by calling `connection.channel` and then use that 50 | channel to interact with the broker. 51 | '''%(options.host, options.vhost) 52 | try: 53 | from IPython import embed 54 | embed(header=hello) 55 | except ImportError: 56 | # Fallback to python console, with benefits 57 | import code 58 | import rlcompleter 59 | import readline 60 | readline.parse_and_bind("tab: complete") 61 | print hello 62 | shell = code.InteractiveConsole( {'Message':Message, 'connection':connection} ) 63 | shell.interact('') 64 | -------------------------------------------------------------------------------- /scripts/rabbit_table_test: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #-*- coding:utf-8 -*- 3 | 4 | ''' 5 | A script for integration testing against rabbit and ensuring that the errata 6 | documented at http://dev.rabbitmq.com/wiki/Amqp091Errata#section_3 is adhered 7 | to. 8 | ''' 9 | 10 | import sys, os 11 | sys.path.append(os.path.abspath(".")) 12 | sys.path.append(os.path.abspath("..")) 13 | 14 | import logging 15 | import random 16 | import socket 17 | from optparse import OptionParser 18 | 19 | from haigha.connection import Connection 20 | from haigha.message import Message 21 | import event 22 | import signal 23 | import time 24 | from decimal import Decimal 25 | from datetime import datetime 26 | 27 | def sigint_cb(*args): 28 | stop_test() 29 | 30 | def stop_test(): 31 | t_end = time.time() 32 | total = 0 33 | logger.info("stopping test") 34 | # Have to iterate on a copy because list will be modified on close() 35 | for channel in channels[:]: 36 | total += channel._count 37 | channel.close() 38 | event.timeout(10, force_quit) 39 | event.timeout(0, check_close_channels) 40 | 41 | def check_close_channels(): 42 | if len(channels): 43 | event.timeout(1, check_close_channels) 44 | else: 45 | disconnect() 46 | 47 | def channel_closed(channel): 48 | channels.remove( channel ) 49 | 50 | def disconnect(): 51 | connection.close() 52 | 53 | def connection_close_cb(): 54 | logger.info("connection closed") 55 | event.abort() 56 | 57 | def force_quit(): 58 | logger.error("force quit!") 59 | event.abort() 60 | 61 | def open_channel(): 62 | channels.append( ChannelTest(connection, 'anexchange') ) 63 | 64 | class ChannelTest: 65 | def __init__(self, connection, exchange): 66 | self._ch = connection.channel() 67 | self._exchange = exchange 68 | self._queue = '%s'%(self._ch.channel_id) 69 | self._count = 0 70 | 71 | self._ch.exchange.declare( self._exchange, 'direct', auto_delete=True ) 72 | self._ch.queue.declare( self._queue, auto_delete=True ) 73 | self._ch.queue.bind( self._queue, self._exchange, self._queue ) 74 | self._ch.basic.consume( self._queue, self._consume ) 75 | 76 | self._headers = { 77 | 'bool' : True, 78 | 'int' : 2**15-1, 79 | 'long' : 2**63-1, 80 | 'float' : 42.55, 81 | 'decimal' : Decimal('-3.14'), 82 | 'str' : 'hello', 83 | 'unicode' : 'Au\xc3\x9ferdem'.decode('utf8'), 84 | 'datetime' : datetime(1985, 10, 25, 4, 20, 20), 85 | 'dict' : {'foo':'bar'}, 86 | 'none' : None, 87 | 'bytearray' : bytearray('\x02\x03'), 88 | 'list' : [], 89 | } 90 | 91 | self._publish() 92 | 93 | def close_and_reopen(self): 94 | if not self._ch.closed: 95 | self.close() 96 | open_channel() 97 | 98 | def close(self): 99 | self._ch.close() 100 | # HACK: Without a proper callback chain, need to delay this so that rabbit 101 | # 2.4.0 can handle the handshake of channel close before we handshake the 102 | # connection close. Otherwise, it gets both close requests in rapid 103 | # succession and doesn't ack either of them, resulting in a force quit 104 | event.timeout( 0.3, channel_closed, self ) 105 | 106 | def _publish(self): 107 | self._count += 1 108 | if not self._ch.closed: 109 | msg = Message( body='hello world', application_headers=self._headers ) 110 | 111 | self._ch.publish( msg, exchange=self._exchange, routing_key=self._queue ) 112 | 113 | def _published(self): 114 | self._publish() 115 | 116 | def _consume(self, msg): 117 | for k,v in self._headers.iteritems(): 118 | if k=='unicode': 119 | v = v.encode('utf8') 120 | if k=='list': 121 | eq = msg.properties['application_headers'][k] is None 122 | else: 123 | eq = msg.properties['application_headers'][k] == v 124 | print k, 'passed? ', eq 125 | stop_test() 126 | 127 | ################################################################### 128 | 129 | parser = OptionParser( 130 | usage='Usage: rabbit_table_test [options]' 131 | ) 132 | parser.add_option('--user', default='guest', type='string') 133 | parser.add_option('--pass', default='guest', dest='password', type='string') 134 | parser.add_option('--vhost', default='/', type='string') 135 | parser.add_option('--host', default='localhost', type='string') 136 | parser.add_option('--debug', default=0, action='count') 137 | 138 | (options,args) = parser.parse_args() 139 | 140 | debug = options.debug 141 | level = logging.DEBUG if debug else logging.INFO 142 | 143 | # Setup logging 144 | logging.basicConfig(level=level, format="[%(levelname)s %(asctime)s] %(message)s" ) 145 | logger = logging.getLogger('haigha') 146 | 147 | channels = [] 148 | 149 | logger.info( 'connecting ...' ) 150 | event.signal( signal.SIGINT, sigint_cb ) 151 | 152 | sock_opts = { 153 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 154 | } 155 | connection = Connection(logger=logger, debug=debug, 156 | user=options.user, password=options.password, 157 | vhost=options.vhost, host=options.host, 158 | heartbeat=None, close_cb=connection_close_cb, 159 | sock_opts=sock_opts) 160 | 161 | open_channel() 162 | 163 | event.dispatch() 164 | -------------------------------------------------------------------------------- /scripts/regression/issue_24: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import sys, os 3 | sys.path.append(os.path.abspath(".")) 4 | 5 | import gevent 6 | import gevent.socket 7 | from gevent.event import AsyncResult 8 | from haigha.connection import Connection 9 | from haigha.message import Message 10 | 11 | # NOTE: Work around a bug in Haigha 0.5.1-0.5.3 that breaks gevent 12 | # compatibility 13 | import haigha 14 | try: 15 | haigha_version = haigha.__version__ 16 | except AttributeError: 17 | pass 18 | else: 19 | from distutils import version 20 | if (version.StrictVersion(haigha_version) >= 21 | version.StrictVersion("0.5.1") 22 | and 23 | version.StrictVersion(haigha_version) <= 24 | version.StrictVersion("0.5.3")): 25 | print >>sys.stderr, \ 26 | "WARNING: DEPLOYING GEVENT-FRIENDLINESS BUG WORK-AROUND FOR HAIGHA v%s" % ( 27 | haigha_version) 28 | from haigha.transports import socket_transport 29 | import gevent.socket 30 | socket_transport.socket = gevent.socket 31 | 32 | 33 | def test_haigha(): 34 | """ 35 | A simple test to check Haigha's connection/channel opening and closing. 36 | 37 | Note that Rabbit MQ must be running 38 | """ 39 | 40 | channel1CloseWaiter = AsyncResult() 41 | connectionCloseWaiter = AsyncResult() 42 | 43 | def handleChannel1Closed(ch, channelImpl): 44 | print "CHANNEL1 CLOSED: %r" % (ch.close_info,) 45 | channel1CloseWaiter.set() 46 | 47 | def handleConnectionClosed(): 48 | print "CONNECTION CLOSED!" 49 | connectionCloseWaiter.set() 50 | 51 | print "Connecting..." 52 | connection = Connection( 53 | user='guest', password='guest', 54 | vhost='/', host='localhost', 55 | heartbeat=None, debug=True, transport="gevent", 56 | close_cb=handleConnectionClosed) 57 | 58 | def readFrames(conn): 59 | while True: 60 | conn.read_frames() 61 | if connectionCloseWaiter.ready(): 62 | break 63 | 64 | # Start Haigha message pump 65 | print "Starting message pump greenlet..." 66 | g = gevent.spawn(readFrames, conn=connection) 67 | 68 | ch1 = connection.channel() 69 | channel1Impl = ch1.channel 70 | ch1.add_close_listener(lambda ch: handleChannel1Closed(ch, channel1Impl)) 71 | 72 | # Close the channels and wait for close-done 73 | print "Closing channel1..." 74 | ch1.close() 75 | with gevent.Timeout(seconds=10) as myTimeout: 76 | channel1CloseWaiter.wait() 77 | 78 | # Close the connection and wait for close-done 79 | print "Closing connection..." 80 | connection.close() 81 | with gevent.Timeout(seconds=10) as myTimeout: 82 | connectionCloseWaiter.wait() 83 | 84 | print "Killing message pump..." 85 | sys.stdout.flush() 86 | 87 | g.kill() 88 | 89 | 90 | 91 | if __name__ == '__main__': 92 | test_haigha() 93 | -------------------------------------------------------------------------------- /scripts/regression/issue_31: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from haigha.connection import Connection 4 | from haigha.message import Message 5 | 6 | connection = Connection(user='guest', password='guest', vhost='/', host='localhost', heartbeat=None, debug=True) 7 | ch = connection.channel() 8 | ch.exchange.declare('test_exchange', 'direct', durable=True) 9 | ch.queue.declare('test_queue', durable=True) 10 | ch.queue.bind('test_queue', 'test_exchange', 'test_key') 11 | ch.basic.publish( Message("hello world", application_headers={'hello':'world'}), 'test_exchange', 'test_key' ) 12 | 13 | connection.close() 14 | 15 | connection = Connection(user='guest', password='guest', vhost='/', host='localhost', heartbeat=None, debug=True) 16 | ch = connection.channel() 17 | ch.exchange.declare('test_exchange', 'direct') 18 | ch.queue.declare('test_queue') 19 | ch.queue.bind('test_queue', 'test_exchange', 'test_key') 20 | ch.basic.publish( Message("hello world", application_headers={'hello':'world'}), 'test_exchange', 'test_key' ) 21 | 22 | connection.close() 23 | -------------------------------------------------------------------------------- /scripts/regression/issue_33: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from haigha.connection import Connection 4 | from haigha.message import Message 5 | 6 | connection = Connection(user='guest', password='guest', vhost='/', host='localhost', heartbeat=None, debug=True) 7 | ch = connection.channel() 8 | ch.exchange.declare('test_exchange', 'direct', durable=True) 9 | ch.queue.declare('test_queue', durable=True) 10 | ch.queue.bind('test_queue', 'test_exchange', 'test_key') 11 | ch.basic.publish( Message(), 'test_exchange', 'test_key' ) 12 | 13 | try: 14 | ch.basic.publish( Message(body=None), 'test_exchange', 'test_key' ) 15 | print 'FAIL did not raise expected type error' 16 | except TypeError: 17 | print 'SUCCESS expected type raised' 18 | 19 | connection.close() 20 | -------------------------------------------------------------------------------- /scripts/stress_test: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #-*- coding:utf-8 -*- 3 | 4 | import sys, os 5 | sys.path.append(os.path.abspath(".")) 6 | sys.path.append(os.path.abspath("..")) 7 | 8 | import logging 9 | import random 10 | import socket 11 | from optparse import OptionParser 12 | 13 | from haigha.connection import Connection 14 | from haigha.message import Message 15 | import event 16 | import signal 17 | import time 18 | import gevent 19 | from gevent import monkey 20 | 21 | def sigint_cb(*args): 22 | logger.info("stopping test") 23 | print_stats() 24 | 25 | # Have to iterate on a copy because list will be modified on close() 26 | for channel in channels[:]: 27 | channel.close() 28 | event.timeout(10, force_quit) 29 | event.timeout(0, check_close_channels) 30 | 31 | def print_stats(): 32 | t_end = time.time() 33 | total = 0 34 | duration = t_end - t_start 35 | 36 | for channel in channels: 37 | total += channel._count 38 | 39 | logger.info("completed %d in %.06f", total, duration) 40 | logger.info("%0.6f msg/s", float(total) / duration ) 41 | logger.info("frames read: %d, %f/s", 42 | connection.frames_read, float(connection.frames_read)/duration ) 43 | logger.info("frames written: %d, %f/s", 44 | connection.frames_written, float(connection.frames_written)/duration ) 45 | 46 | def check_close_channels(): 47 | if len(channels): 48 | event.timeout(1, check_close_channels) 49 | else: 50 | disconnect() 51 | 52 | def channel_closed(channel): 53 | channels.remove( channel ) 54 | 55 | def disconnect(): 56 | connection.close() 57 | 58 | def connection_close_cb(): 59 | logger.info("connection closed") 60 | event.abort() 61 | 62 | def force_quit(): 63 | logger.error("force quit!") 64 | event.abort() 65 | 66 | def open_channel(): 67 | channels.append( ChannelTest(connection, random.choice(exchanges)) ) 68 | 69 | class ChannelTest: 70 | def __init__(self, connection, exchange): 71 | self._ch = connection.channel() 72 | self._exchange = exchange 73 | self._queue = '%s'%(self._ch.channel_id) 74 | self._count = 0 75 | 76 | self._ch.exchange.declare( self._exchange, 'direct', auto_delete=True ) 77 | self._ch.queue.declare( self._queue, auto_delete=True ) 78 | self._ch.queue.bind( self._queue, self._exchange, self._queue ) 79 | self._ch.basic.consume( self._queue, self._consume ) 80 | 81 | #self._publish() 82 | event.timeout( 1, self._publish ) 83 | 84 | @property 85 | def _transactions(self): 86 | return options.tx 87 | 88 | def close_and_reopen(self): 89 | if not self._ch.closed: 90 | self.close() 91 | open_channel() 92 | 93 | def close(self): 94 | self._ch.close() 95 | # HACK: Without a proper callback chain, need to delay this so that rabbit 96 | # 2.4.0 can handle the handshake of channel close before we handshake the 97 | # connection close. Otherwise, it gets both close requests in rapid 98 | # succession and doesn't ack either of them, resulting in a force quit 99 | #event.timeout( 0.3, channel_closed, self ) 100 | channel_closed(self) 101 | 102 | def _publish(self): 103 | self._count += 1 104 | if not self._ch.closed: 105 | msg = Message( body='hello world' ) 106 | 107 | if self._transactions: 108 | self._ch.publish_synchronous( msg, exchange=self._exchange, routing_key=self._queue, cb=self._published ) 109 | else: 110 | self._ch.publish( msg, exchange=self._exchange, routing_key=self._queue ) 111 | 112 | def _published(self): 113 | self._publish() 114 | 115 | def _consume(self, msg): 116 | if not self._transactions: 117 | self._publish() 118 | 119 | ################################################################### 120 | 121 | parser = OptionParser( 122 | usage='Usage: stress_test [options]' 123 | ) 124 | parser.add_option('--user', default='guest', type='string') 125 | parser.add_option('--pass', default='guest', dest='password', type='string') 126 | parser.add_option('--vhost', default='/', type='string') 127 | parser.add_option('--host', default='localhost', type='string') 128 | parser.add_option('--tx', default=False, action='store_true' ) 129 | parser.add_option('--profile', default=False, action='store_true' ) 130 | parser.add_option('--channels', default=500, type='int') 131 | parser.add_option('--debug', default=0, action='count') 132 | parser.add_option('--time', default=0, type='int') 133 | parser.add_option('--transport', default='event', choices=['event', 'gevent', 'gevent_pool']) 134 | 135 | (options,args) = parser.parse_args() 136 | 137 | debug = options.debug 138 | level = logging.DEBUG if debug else logging.INFO 139 | 140 | # Setup logging 141 | logging.basicConfig(level=level, format="[%(levelname)s %(asctime)s] %(message)s" ) 142 | logger = logging.getLogger('haigha') 143 | 144 | channels = [] 145 | 146 | logger.info( 'connecting with transport %s ...', options.transport ) 147 | 148 | # Need to monkey patch before the connection is made 149 | if options.transport in ('gevent', 'gevent_pool'): 150 | monkey.patch_all() 151 | 152 | sock_opts = { 153 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 154 | } 155 | connection = Connection(logger=logger, debug=debug, 156 | user=options.user, password=options.password, 157 | vhost=options.vhost, host=options.host, 158 | heartbeat=None, close_cb=connection_close_cb, 159 | sock_opts=sock_opts, 160 | transport=options.transport) 161 | 162 | exchanges = ['publish-%d'%(x) for x in xrange(0,10)] 163 | 164 | for x in xrange(0,options.channels): 165 | open_channel() 166 | 167 | if options.transport=='event': 168 | if options.time: 169 | event.timeout( options.time, sigint_cb ) 170 | 171 | t_start = time.time() 172 | event.signal( signal.SIGINT, sigint_cb ) 173 | 174 | if options.profile: 175 | import cProfile 176 | cProfile.run( 'event.dispatch()', 'profile.pstats' ) 177 | else: 178 | event.dispatch() 179 | 180 | elif options.transport=='gevent': 181 | def frame_loop(): 182 | while True: 183 | connection.read_frames() 184 | gevent.sleep(0) 185 | 186 | t_start = time.time() 187 | greenlet = gevent.spawn( frame_loop ) 188 | greenlet.join( options.time ) 189 | 190 | print_stats() 191 | 192 | elif options.transport=='gevent_pool': 193 | def frame_loop(): 194 | while True: 195 | connection.read_frames() 196 | gevent.sleep(0) 197 | 198 | t_start = time.time() 199 | connection.transport.pool.spawn( frame_loop ) 200 | connection.transport.pool.join( options.time ) 201 | 202 | print_stats() 203 | -------------------------------------------------------------------------------- /scripts/test: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | #-*- coding:utf-8 -*- 3 | import unittest2 as unittest 4 | import os 5 | import sys 6 | 7 | project_root = os.path.split(os.path.abspath(os.path.dirname(__file__)))[0] 8 | 9 | loader = unittest.TestLoader() 10 | suite = loader.discover('haigha.tests', pattern='*_test.py', top_level_dir=project_root) 11 | 12 | unittest.TextTestRunner(verbosity=1).run(suite) 13 | 14 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import haigha 2 | import os 3 | 4 | try: 5 | from setuptools import setup 6 | except ImportError: 7 | from distutils.core import setup 8 | 9 | 10 | requirements = map(str.strip, open('requirements.txt').readlines()) 11 | 12 | setup( 13 | name='haigha', 14 | version=haigha.__version__, 15 | author='Vitaly Babiy, Aaron Westendorf', 16 | author_email="vbabiy@agoragames.com, aaron@agoragames.com", 17 | packages = ['haigha', 'haigha.frames', 'haigha.classes', 'haigha.transports', 'haigha.connections'], 18 | install_requires = requirements, 19 | url='https://github.com/agoragames/haigha', 20 | license="LICENSE.txt", 21 | description='Synchronous and asynchronous AMQP client library', 22 | long_description=open('README.rst').read(), 23 | keywords=['python', 'amqp', 'event', 'rabbitmq'], 24 | classifiers=[ 25 | 'Development Status :: 5 - Production/Stable', 26 | 'License :: OSI Approved :: BSD License', 27 | "Intended Audience :: Developers", 28 | "Operating System :: POSIX", 29 | "Topic :: Communications", 30 | "Topic :: System :: Distributed Computing", 31 | "Topic :: Software Development :: Libraries :: Python Modules", 32 | 'Programming Language :: Python', 33 | 'Programming Language :: Python :: 2.7', 34 | #'Programming Language :: Python :: 3', 35 | 'Topic :: Software Development :: Libraries' 36 | ] 37 | ) 38 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/tests/__init__.py -------------------------------------------------------------------------------- /tests/integration/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/tests/integration/__init__.py -------------------------------------------------------------------------------- /tests/integration/channel_basic_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | # 8 | # Integration tests for Channel.basic 9 | # 10 | 11 | 12 | # Disable "no member" pylint error since Haigha's channel class members get 13 | # added at runtime. 14 | # 15 | # e.g., "E1103: Instance of 'Channel' has no 'exchange' member (but some types 16 | # could not be inferred)" 17 | # 18 | # pylint: disable=E1103 19 | 20 | 21 | import logging 22 | import socket 23 | import unittest 24 | 25 | 26 | from haigha.channel import Channel 27 | from haigha.connection import Connection 28 | from haigha.message import Message 29 | 30 | 31 | class TestOptions(object): # pylint: disable=R0903 32 | '''Configuration settings''' 33 | user = 'guest' 34 | password = 'guest' 35 | vhost = '/' 36 | host = 'localhost' 37 | debug = False 38 | 39 | 40 | _OPTIONS = TestOptions() 41 | _LOG = None 42 | 43 | 44 | def setUpModule(): # pylint: disable=C0103 45 | '''Unittest fixture for module-level initialization''' 46 | 47 | global _LOG # pylint: disable=W0603 48 | 49 | 50 | # Setup logging 51 | log_level = logging.DEBUG if _OPTIONS.debug else logging.INFO 52 | logging.basicConfig(level=log_level, 53 | format="[%(levelname)s %(asctime)s] %(message)s") 54 | _LOG = logging.getLogger('haigha') 55 | 56 | 57 | class _CallbackSink(object): 58 | '''Callback sink; an instance of this class may be passed as a callback 59 | and it will store the callback args in the values instance attribute 60 | ''' 61 | 62 | __slots__ = ('values',) 63 | 64 | def __init__(self): 65 | self.values = None 66 | self.reset() 67 | 68 | def reset(self): 69 | '''Reset the args buffer''' 70 | self.values = [] 71 | 72 | def __repr__(self): 73 | return "%s(ready=%s, values=%.255r)" % (self.__class__.__name__, 74 | self.ready, 75 | self.values) 76 | 77 | def __call__(self, *args): 78 | self.values.append(args) 79 | 80 | @property 81 | def ready(self): 82 | '''True if called; False if not called''' 83 | return bool(self.values) 84 | 85 | 86 | class ChannelBasicTests(unittest.TestCase): 87 | '''Integration tests for Channel.basic''' 88 | 89 | def _connect_to_broker(self): 90 | ''' Connect to broker and regisiter cleanup action to disconnect 91 | 92 | :returns: connection instance 93 | :rtype: `haigha.connection.Connection` 94 | ''' 95 | sock_opts = { 96 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 97 | } 98 | connection = Connection( 99 | logger=_LOG, 100 | debug=_OPTIONS.debug, 101 | user=_OPTIONS.user, 102 | password=_OPTIONS.password, 103 | vhost=_OPTIONS.vhost, 104 | host=_OPTIONS.host, 105 | heartbeat=None, 106 | sock_opts=sock_opts, 107 | transport='socket') 108 | self.addCleanup(lambda: connection.close(disconnect=True) 109 | if not connection.closed else None) 110 | 111 | return connection 112 | 113 | def test_unroutable_message_is_returned(self): 114 | connection = self._connect_to_broker() 115 | 116 | ch = connection.channel() 117 | self.addCleanup(ch.close) 118 | 119 | _LOG.info('Declaring exchange "foo"') 120 | ch.exchange.declare('foo', 'direct') 121 | self.addCleanup(ch.exchange.delete, 'foo') 122 | 123 | callback_sink = _CallbackSink() 124 | ch.basic.set_return_listener(callback_sink) 125 | 126 | _LOG.info( 127 | 'Publishing to exchange "foo" on route "nullroute" mandatory=True') 128 | mid = ch.basic.publish(Message('hello world'), 'foo', 'nullroute', 129 | mandatory=True) 130 | _LOG.info('Published message mid %s to foo/nullroute', mid) 131 | 132 | # Wait for return of unroutable message 133 | while not callback_sink.ready: 134 | connection.read_frames() 135 | 136 | # Validate returned message 137 | ((msg,),) = callback_sink.values 138 | self.assertEqual(msg.body, 'hello world') 139 | 140 | self.assertIsNone(msg.delivery_info) 141 | self.assertIsNotNone(msg.return_info) 142 | 143 | return_info = msg.return_info 144 | self.assertItemsEqual( 145 | ['channel', 'reply_code', 'reply_text', 'exchange', 'routing_key'], 146 | return_info.keys()) 147 | self.assertIs(return_info['channel'], ch) 148 | self.assertEqual(return_info['reply_code'], 312) 149 | self.assertEqual(return_info['reply_text'], 'NO_ROUTE') 150 | self.assertEqual(return_info['exchange'], 'foo') 151 | self.assertEqual(return_info['routing_key'], 'nullroute') 152 | -------------------------------------------------------------------------------- /tests/integration/rabbit_extensions_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | # 8 | # Integration tests for RabbitMQ extensions 9 | # 10 | 11 | 12 | # Disable "no member" pylint error since Haigha's channel class members get 13 | # added at runtime. 14 | # 15 | # e.g., "E1103: Instance of 'Channel' has no 'exchange' member (but some types 16 | # could not be inferred)" 17 | # 18 | # pylint: disable=E1103 19 | 20 | 21 | # Disable pylint warning regarding protected member access 22 | # 23 | # pylint: disable=W0212 24 | 25 | 26 | # Disable pylint notification about missing method docstring 27 | # 28 | # pylint: disable=C0111 29 | 30 | 31 | 32 | import logging 33 | import socket 34 | import unittest 35 | 36 | from haigha.connections.rabbit_connection import RabbitConnection 37 | from haigha.message import Message 38 | 39 | 40 | class TestOptions(object): # pylint: disable=R0903 41 | '''Configuration settings''' 42 | user = 'guest' 43 | password = 'guest' 44 | vhost = '/' 45 | host = 'localhost' 46 | debug = False 47 | 48 | 49 | _OPTIONS = TestOptions() 50 | _LOG = None 51 | 52 | 53 | def setUpModule(): # pylint: disable=C0103 54 | '''Unittest fixture for module-level initialization''' 55 | 56 | global _LOG # pylint: disable=W0603 57 | 58 | 59 | # Setup logging 60 | log_level = logging.DEBUG if _OPTIONS.debug else logging.INFO 61 | logging.basicConfig(level=log_level, 62 | format="[%(levelname)s %(asctime)s] %(message)s") 63 | _LOG = logging.getLogger('haigha') 64 | 65 | 66 | class _CallbackSink(object): 67 | '''Callback sink; an instance of this class may be passed as a callback 68 | and it will store the callback args in the values instance attribute 69 | ''' 70 | 71 | __slots__ = ('values',) 72 | 73 | def __init__(self): 74 | self.values = None 75 | self.reset() 76 | 77 | def reset(self): 78 | '''Reset the args buffer''' 79 | self.values = [] 80 | 81 | def __repr__(self): 82 | return "%s(ready=%s, values=%.255r)" % (self.__class__.__name__, 83 | self.ready, 84 | self.values) 85 | 86 | def __call__(self, *args): 87 | self.values.append(args) 88 | 89 | @property 90 | def ready(self): 91 | '''True if called; False if not called''' 92 | return bool(self.values) 93 | 94 | 95 | class _PubackState(_CallbackSink): 96 | '''An instance of this class acts as a context manager that registers for 97 | basic.ack/nack notificagtions and records ACK or NACK callback 98 | ''' 99 | 100 | __slots__ = ('_channel') 101 | 102 | ACK = 1 103 | NACK = 2 104 | 105 | def __init__(self, channel): 106 | self._channel = channel 107 | super(_PubackState, self).__init__() 108 | 109 | def __enter__(self): 110 | self._channel.basic.set_ack_listener(self.handle_ack) 111 | self._channel.basic.set_nack_listener(self.handle_nack) 112 | return self 113 | 114 | def __exit__(self, *_args): 115 | self._channel.basic.set_ack_listener(None) 116 | self._channel.basic.set_nack_listener(None) 117 | self.reset() 118 | self._channel = None 119 | 120 | def handle_ack(self, delivery_tag): 121 | '''Message Ack'ed in RabbitMQ Publisher Acknowledgments mode''' 122 | _LOG.debug("Message ACKed: tag=%s", delivery_tag) 123 | 124 | assert not self.ready, (delivery_tag, self.values) 125 | 126 | self(self.ACK, delivery_tag) 127 | 128 | def handle_nack(self, delivery_tag): 129 | '''Message Nack'ed in RabbitMQ Publisher Acknowledgments mode''' 130 | _LOG.error("Message NACKed: tag=%s", delivery_tag) 131 | 132 | assert not self.ready, (delivery_tag, self.values) 133 | 134 | self(self.NACK, delivery_tag) 135 | 136 | 137 | class RabbitExtensionsTests(unittest.TestCase): 138 | '''Integration tests for RabbitMQ-specific extensions''' 139 | 140 | def _connect_to_broker(self): 141 | ''' Connect to broker and regisiter cleanup action to disconnect 142 | 143 | :returns: connection instance 144 | :rtype: `haigha.connections.rabbit_connection.Connection` 145 | ''' 146 | sock_opts = { 147 | (socket.IPPROTO_TCP, socket.TCP_NODELAY) : 1, 148 | } 149 | connection = RabbitConnection( 150 | logger=_LOG, 151 | debug=_OPTIONS.debug, 152 | user=_OPTIONS.user, 153 | password=_OPTIONS.password, 154 | vhost=_OPTIONS.vhost, 155 | host=_OPTIONS.host, 156 | heartbeat=None, 157 | sock_opts=sock_opts, 158 | transport='socket') 159 | self.addCleanup(lambda: connection.close(disconnect=True) 160 | if not connection.closed else None) 161 | 162 | return connection 163 | 164 | def test_puback_and_internal_exchange(self): 165 | # NOTE: this test method was created from the old rabbit_extensions 166 | # script contents and enhanced with ack validation and waiting for get 167 | # to complete. It needs some more work, including splitting into 168 | # multiple tests: sock opts, exchange-to-exchange binding, and 169 | # publisher-acks. 170 | connection = self._connect_to_broker() 171 | 172 | ch = connection.channel() 173 | self.addCleanup(ch.close) 174 | 175 | _LOG.info('Declaring exchange "foo"') 176 | ch.exchange.declare('foo', 'direct', auto_delete=True) 177 | self.addCleanup(ch.exchange.delete, 'foo') 178 | 179 | _LOG.info('Declaring internal exchange "fooint"') 180 | ch.exchange.declare('fooint', 'direct', internal=True, auto_delete=True, 181 | arguments={}) 182 | self.addCleanup(ch.exchange.delete, 'fooint') 183 | 184 | _LOG.info('Binding "fooint" to "foo" on route "route"') 185 | ch.exchange.bind('fooint', 'foo', 'route') 186 | 187 | _LOG.info( 188 | 'Binding queue "bar" to exchange "fooint" on route "route"') 189 | ch.queue.declare('bar', auto_delete=True) 190 | self.addCleanup(ch.queue.delete, 'bar') 191 | ch.queue.bind('bar', 'fooint', 'route') 192 | 193 | _LOG.info('Enabling publisher confirmations') 194 | ch.confirm.select() 195 | 196 | _LOG.info('Publishing to exchange "foo" on route "route"') 197 | 198 | with _PubackState(ch) as puback_state: 199 | mid = ch.basic.publish(Message('hello world'), 'foo', 'route') 200 | _LOG.info('Published message mid %s', mid) 201 | while not puback_state.ready: 202 | connection.read_frames() 203 | 204 | ((how, response_tag),) = puback_state.values 205 | self.assertEqual(how, puback_state.ACK) 206 | self.assertEqual(response_tag, mid) 207 | 208 | consumer = _CallbackSink() 209 | msg = ch.basic.get('bar', consumer=consumer) 210 | _LOG.info('GET %s', msg) 211 | # Wait for the rest of the message body to come in, not just the 212 | # Basic.Get-ok frame 213 | while not consumer.ready: 214 | connection.read_frames() 215 | ((msg,),) = consumer.values 216 | _LOG.info('GOT %s', msg) 217 | self.assertEqual(msg.body, 'hello world') 218 | 219 | 220 | _LOG.info( 221 | 'Publishing to exchange "foo" on route "nullroute" mandatory=False') 222 | with _PubackState(ch) as puback_state: 223 | mid = ch.basic.publish(Message('hello world'), 'foo', 'nullroute') 224 | _LOG.info('Published message mid %s to foo/nullroute', mid) 225 | while not puback_state.ready: 226 | connection.read_frames() 227 | 228 | ((how, response_tag),) = puback_state.values 229 | self.assertEqual(how, puback_state.ACK) 230 | self.assertEqual(response_tag, mid) 231 | 232 | def test_unroutable_message_is_returned_with_puback(self): 233 | connection = self._connect_to_broker() 234 | 235 | ch = connection.channel() 236 | self.addCleanup(ch.close) 237 | 238 | _LOG.info('Declaring exchange "foo"') 239 | ch.exchange.declare('foo', 'direct') 240 | self.addCleanup(ch.exchange.delete, 'foo') 241 | 242 | callback_sink = _CallbackSink() 243 | ch.basic.set_return_listener(callback_sink) 244 | 245 | _LOG.info('Enabling publisher confirmations') 246 | ch.confirm.select() 247 | 248 | _LOG.info( 249 | 'Publishing to exchange "foo" on route "nullroute" mandatory=True') 250 | with _PubackState(ch) as puback_state: 251 | mid = ch.basic.publish(Message('hello world'), 'foo', 'nullroute', 252 | mandatory=True) 253 | _LOG.info('Published message mid %s to foo/nullroute', mid) 254 | 255 | # Wait for ack 256 | while not puback_state.ready: 257 | connection.read_frames() 258 | 259 | ((how, response_tag),) = puback_state.values 260 | self.assertEqual(how, puback_state.ACK) 261 | self.assertEqual(response_tag, mid) 262 | 263 | 264 | # Verify that unroutable message was returned 265 | 266 | self.assertEqual(len(callback_sink.values), 1) 267 | 268 | ((msg,),) = callback_sink.values 269 | self.assertEqual(msg.body, 'hello world') 270 | 271 | self.assertIsNone(msg.delivery_info) 272 | self.assertIsNotNone(msg.return_info) 273 | 274 | return_info = msg.return_info 275 | self.assertItemsEqual( 276 | ['channel', 'reply_code', 'reply_text', 'exchange', 'routing_key'], 277 | return_info.keys()) 278 | self.assertIs(return_info['channel'], ch) 279 | self.assertEqual(return_info['reply_code'], 312) 280 | self.assertEqual(return_info['reply_text'], 'NO_ROUTE') 281 | self.assertEqual(return_info['exchange'], 'foo') 282 | self.assertEqual(return_info['routing_key'], 'nullroute') 283 | 284 | def test_basic_cancel_from_broker(self): 285 | connection = self._connect_to_broker() 286 | 287 | ch = connection.channel() 288 | self.addCleanup(ch.close) 289 | 290 | queue, _, _ = ch.queue.declare('', auto_delete=True) 291 | 292 | # Start consumer 293 | cancel_cb = _CallbackSink() 294 | consumer = _CallbackSink() 295 | consumer_tag = ch.basic._generate_consumer_tag() 296 | ch.basic.consume(queue, consumer=consumer, consumer_tag=consumer_tag, 297 | cancel_cb=cancel_cb) 298 | 299 | # Delete the queue to force consumer cancellation 300 | ch.queue.delete(queue) 301 | 302 | # Wait for Basic.Cancel from server 303 | while not cancel_cb.ready: 304 | connection.read_frames() 305 | 306 | ((rx_consumer_tag,),) = cancel_cb.values 307 | self.assertEqual(rx_consumer_tag, consumer_tag) 308 | 309 | 310 | if __name__ == '__main__': 311 | unittest.main() 312 | -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/tests/unit/__init__.py -------------------------------------------------------------------------------- /tests/unit/channel_pool_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from collections import deque 8 | 9 | from chai import Chai 10 | 11 | from haigha.channel_pool import ChannelPool 12 | 13 | 14 | class ChannelPoolTest(Chai): 15 | 16 | def test_init(self): 17 | c = ChannelPool('connection') 18 | self.assertEquals('connection', c._connection) 19 | self.assertEquals(set(), c._free_channels) 20 | assert_equals(None, c._size) 21 | assert_equals(0, c._channels) 22 | assert_equals(deque(), c._queue) 23 | 24 | c = ChannelPool('connection', size=50) 25 | self.assertEquals('connection', c._connection) 26 | self.assertEquals(set(), c._free_channels) 27 | assert_equals(50, c._size) 28 | assert_equals(0, c._channels) 29 | assert_equals(deque(), c._queue) 30 | 31 | def test_publish_without_user_cb(self): 32 | ch = mock() 33 | cp = ChannelPool(None) 34 | 35 | expect(cp._get_channel).returns(ch) 36 | expect(ch.publish_synchronous).args( 37 | 'arg1', 'arg2', cb=var('cb'), doit='harder') 38 | 39 | cp.publish('arg1', 'arg2', doit='harder') 40 | assert_equals(set(), cp._free_channels) 41 | 42 | # run committed callback 43 | var('cb').value() 44 | assert_equals(set([ch]), cp._free_channels) 45 | 46 | def test_publish_with_user_cb(self): 47 | ch = mock() 48 | cp = ChannelPool(None) 49 | user_cb = mock() 50 | 51 | expect(cp._get_channel).returns(ch) 52 | expect(ch.publish_synchronous).args( 53 | 'arg1', 'arg2', cb=var('cb'), doit='harder') 54 | 55 | cp.publish('arg1', 'arg2', cb=user_cb, doit='harder') 56 | assert_equals(set(), cp._free_channels) 57 | 58 | expect(user_cb) 59 | var('cb').value() 60 | assert_equals(set([ch]), cp._free_channels) 61 | 62 | def test_publish_resends_queued_messages_if_channel_is_active(self): 63 | ch = mock() 64 | cp = ChannelPool(None) 65 | user_cb = mock() 66 | ch.active = True 67 | ch.closed = False 68 | cp._queue.append((('a1', 'a2'), {'cb': 'foo', 'yo': 'dawg'})) 69 | 70 | expect(cp._get_channel).returns(ch) 71 | expect(ch.publish_synchronous).args( 72 | 'arg1', 'arg2', cb=var('cb'), doit='harder') 73 | 74 | cp.publish('arg1', 'arg2', cb=user_cb, doit='harder') 75 | assert_equals(set(), cp._free_channels) 76 | assert_equals(1, len(cp._queue)) 77 | 78 | expect(cp._process_queue) 79 | expect(user_cb) 80 | var('cb').value() 81 | assert_equals(set([ch]), cp._free_channels) 82 | 83 | def test_publish_does_not_resend_queued_messages_if_channel_is_inactive(self): 84 | ch = mock() 85 | cp = ChannelPool(None) 86 | user_cb = mock() 87 | ch.active = True 88 | ch.closed = False 89 | cp._queue.append((('a1', 'a2'), {'cb': 'foo', 'yo': 'dawg'})) 90 | 91 | expect(cp._get_channel).returns(ch) 92 | expect(ch.publish_synchronous).args( 93 | 'arg1', 'arg2', cb=var('cb'), doit='harder') 94 | 95 | cp.publish('arg1', 'arg2', cb=user_cb, doit='harder') 96 | assert_equals(set(), cp._free_channels) 97 | assert_equals(1, len(cp._queue)) 98 | 99 | ch.active = False 100 | 101 | stub(cp._process_queue) 102 | expect(user_cb) 103 | var('cb').value() 104 | assert_equals(set([ch]), cp._free_channels) 105 | assert_equals(1, len(cp._queue)) 106 | 107 | def test_publish_does_not_resend_queued_messages_if_channel_is_closed(self): 108 | ch = mock() 109 | cp = ChannelPool(None) 110 | user_cb = mock() 111 | ch.active = True 112 | ch.closed = False 113 | cp._queue.append((('a1', 'a2'), {'cb': 'foo', 'yo': 'dawg'})) 114 | 115 | expect(cp._get_channel).returns(ch) 116 | expect(ch.publish_synchronous).args( 117 | 'arg1', 'arg2', cb=var('cb'), doit='harder') 118 | 119 | cp.publish('arg1', 'arg2', cb=user_cb, doit='harder') 120 | assert_equals(set(), cp._free_channels) 121 | assert_equals(1, len(cp._queue)) 122 | 123 | ch.closed = True 124 | stub(cp._process_queue) 125 | expect(user_cb) 126 | var('cb').value() 127 | assert_equals(set([ch]), cp._free_channels) 128 | assert_equals(1, len(cp._queue)) 129 | 130 | def test_publish_searches_for_active_channel(self): 131 | ch1 = mock() 132 | ch2 = mock() 133 | ch3 = mock() 134 | ch1.active = ch2.active = False 135 | ch3.active = True 136 | cp = ChannelPool(None) 137 | 138 | expect(cp._get_channel).returns(ch1) 139 | expect(cp._get_channel).returns(ch2) 140 | expect(cp._get_channel).returns(ch3) 141 | expect(ch3.publish_synchronous).args('arg1', 'arg2', cb=ignore()) 142 | 143 | cp.publish('arg1', 'arg2') 144 | self.assertEquals(set([ch1, ch2]), cp._free_channels) 145 | 146 | def test_publish_appends_to_queue_when_no_ready_channels(self): 147 | cp = ChannelPool(None) 148 | 149 | expect(cp._get_channel).returns(None) 150 | 151 | cp.publish('arg1', 'arg2', arg3='foo', cb='usercb') 152 | self.assertEquals(set(), cp._free_channels) 153 | assert_equals(deque([(('arg1', 'arg2'), {'arg3': 'foo', 'cb': 'usercb'})]), 154 | cp._queue) 155 | 156 | def test_publish_appends_to_queue_when_no_ready_channels_out_of_several(self): 157 | ch1 = mock() 158 | cp = ChannelPool(None) 159 | ch1.active = False 160 | 161 | expect(cp._get_channel).returns(ch1) 162 | expect(cp._get_channel).returns(None) 163 | 164 | cp.publish('arg1', 'arg2', arg3='foo', cb='usercb') 165 | self.assertEquals(set([ch1]), cp._free_channels) 166 | assert_equals(deque([(('arg1', 'arg2'), {'arg3': 'foo', 'cb': 'usercb'})]), 167 | cp._queue) 168 | 169 | def test_process_queue(self): 170 | cp = ChannelPool(None) 171 | cp._queue = deque([ 172 | (('foo',), {'a': 1}), 173 | (('bar',), {'b': 2}), 174 | ]) 175 | expect(cp.publish).args('foo', a=1) 176 | expect(cp.publish).args('bar', b=2) 177 | 178 | cp._process_queue() 179 | cp._process_queue() 180 | cp._process_queue() 181 | 182 | def test_get_channel_returns_new_when_none_free_and_not_at_limit(self): 183 | conn = mock() 184 | cp = ChannelPool(conn) 185 | cp._channels = 1 186 | 187 | with expect(conn.channel).returns(mock()) as newchannel: 188 | expect(newchannel.add_close_listener).args(cp._channel_closed_cb) 189 | self.assertEquals(newchannel, cp._get_channel()) 190 | self.assertEquals(set(), cp._free_channels) 191 | assert_equals(2, cp._channels) 192 | 193 | def test_get_channel_returns_new_when_none_free_and_at_limit(self): 194 | conn = mock() 195 | cp = ChannelPool(conn, 1) 196 | cp._channels = 1 197 | 198 | stub(conn.channel) 199 | 200 | self.assertEquals(None, cp._get_channel()) 201 | self.assertEquals(set(), cp._free_channels) 202 | 203 | def test_get_channel_when_one_free_and_not_closed(self): 204 | conn = mock() 205 | ch = mock() 206 | ch.closed = False 207 | cp = ChannelPool(conn) 208 | cp._free_channels = set([ch]) 209 | 210 | self.assertEquals(ch, cp._get_channel()) 211 | self.assertEquals(set(), cp._free_channels) 212 | 213 | def test_get_channel_when_two_free_and_one_closed(self): 214 | # Because we can't mock builtins .... 215 | class Set(set): 216 | 217 | def pop(self): 218 | pass 219 | 220 | conn = mock() 221 | ch1 = mock() 222 | ch1.closed = True 223 | ch2 = mock() 224 | ch2.closed = False 225 | cp = ChannelPool(conn) 226 | cp._free_channels = Set([ch1, ch2]) 227 | cp._channels = 2 228 | 229 | # Because we want them in order 230 | expect(cp._free_channels.pop).returns( 231 | ch1).side_effect(super(Set, cp._free_channels).pop) 232 | expect(cp._free_channels.pop).returns( 233 | ch2).side_effect(super(Set, cp._free_channels).pop) 234 | 235 | self.assertEquals(ch2, cp._get_channel()) 236 | self.assertEquals(set(), cp._free_channels) 237 | assert_equals(2, cp._channels) 238 | 239 | def test_get_channel_when_two_free_and_all_closed(self): 240 | conn = mock() 241 | ch1 = mock() 242 | ch1.closed = True 243 | ch2 = mock() 244 | ch2.closed = True 245 | cp = ChannelPool(conn) 246 | cp._free_channels = set([ch1, ch2]) 247 | cp._channels = 2 248 | 249 | with expect(conn.channel).returns(mock()) as newchannel: 250 | expect(newchannel.add_close_listener).args(cp._channel_closed_cb) 251 | self.assertEquals(newchannel, cp._get_channel()) 252 | 253 | self.assertEquals(set(), cp._free_channels) 254 | assert_equals(3, cp._channels) 255 | 256 | def test_channel_closed_cb(self): 257 | cp = ChannelPool(None) 258 | cp._channels = 32 259 | 260 | expect(cp._process_queue) 261 | cp._channel_closed_cb('channel') 262 | assert_equals(31, cp._channels) 263 | -------------------------------------------------------------------------------- /tests/unit/classes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/tests/unit/classes/__init__.py -------------------------------------------------------------------------------- /tests/unit/classes/channel_class_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.channel import Channel 10 | from haigha.classes import channel_class 11 | from haigha.classes.protocol_class import ProtocolClass 12 | from haigha.classes.channel_class import ChannelClass 13 | from haigha.frames.method_frame import MethodFrame 14 | from haigha.writer import Writer 15 | 16 | 17 | class ChannelClassTest(Chai): 18 | 19 | def setUp(self): 20 | super(ChannelClassTest, self).setUp() 21 | connection = mock() 22 | ch = Channel(connection, 42, {}) 23 | connection._logger = mock() 24 | self.klass = ChannelClass(ch) 25 | 26 | def test_init(self): 27 | expect(ProtocolClass.__init__).args('foo', a='b') 28 | 29 | klass = ChannelClass.__new__(ChannelClass) 30 | klass.__init__('foo', a='b') 31 | 32 | assert_equals( 33 | { 34 | 11: klass._recv_open_ok, 35 | 20: klass._recv_flow, 36 | 21: klass._recv_flow_ok, 37 | 40: klass._recv_close, 38 | 41: klass._recv_close_ok, 39 | }, klass.dispatch_map) 40 | assert_equals(None, klass._flow_control_cb) 41 | 42 | def test_cleanup(self): 43 | self.klass._cleanup() 44 | assert_equals(None, self.klass._channel) 45 | assert_equals(None, self.klass.dispatch_map) 46 | 47 | def test_set_flow_cb(self): 48 | assert_equals(None, self.klass._flow_control_cb) 49 | self.klass.set_flow_cb('foo') 50 | assert_equals('foo', self.klass._flow_control_cb) 51 | 52 | def test_open(self): 53 | writer = mock() 54 | expect(mock(channel_class, 'Writer')).returns(writer) 55 | expect(writer.write_shortstr).args('') 56 | expect(mock(channel_class, 'MethodFrame')).args( 57 | 42, 20, 10, writer).returns('frame') 58 | expect(self.klass.send_frame).args('frame') 59 | expect(self.klass.channel.add_synchronous_cb).args( 60 | self.klass._recv_open_ok) 61 | 62 | self.klass.open() 63 | 64 | def test_recv_open_ok(self): 65 | expect(self.klass.channel._notify_open_listeners) 66 | self.klass._recv_open_ok('methodframe') 67 | 68 | def test_activate_when_not_active(self): 69 | self.klass.channel._active = False 70 | expect(self.klass._send_flow).args(True) 71 | self.klass.activate() 72 | 73 | def test_activate_when_active(self): 74 | self.klass.channel._active = True 75 | stub(self.klass._send_flow) 76 | self.klass.activate() 77 | 78 | def test_deactivate_when_not_active(self): 79 | self.klass.channel._active = False 80 | stub(self.klass._send_flow) 81 | self.klass.deactivate() 82 | 83 | def test_deactivate_when_active(self): 84 | self.klass.channel._active = True 85 | expect(self.klass._send_flow).args(False) 86 | self.klass.deactivate() 87 | 88 | def test_send_flow(self): 89 | writer = mock() 90 | expect(mock(channel_class, 'Writer')).returns(writer) 91 | expect(writer.write_bit).args('active') 92 | expect(mock(channel_class, 'MethodFrame')).args( 93 | 42, 20, 20, writer).returns('frame') 94 | expect(self.klass.send_frame).args('frame') 95 | expect(self.klass.channel.add_synchronous_cb).args( 96 | self.klass._recv_flow_ok) 97 | 98 | self.klass._send_flow('active') 99 | 100 | def test_recv_flow_no_cb(self): 101 | self.klass._flow_control_cb = None 102 | rframe = mock() 103 | writer = mock() 104 | expect(rframe.args.read_bit).returns('active') 105 | 106 | expect(mock(channel_class, 'Writer')).returns(writer) 107 | expect(writer.write_bit).args('active') 108 | expect(mock(channel_class, 'MethodFrame')).args( 109 | 42, 20, 21, writer).returns('frame') 110 | expect(self.klass.send_frame).args('frame') 111 | 112 | self.klass._recv_flow(rframe) 113 | assert_equals('active', self.klass.channel._active) 114 | 115 | def test_recv_flow_with_cb(self): 116 | self.klass._flow_control_cb = mock() 117 | rframe = mock() 118 | writer = mock() 119 | expect(rframe.args.read_bit).returns('active') 120 | 121 | expect(mock(channel_class, 'Writer')).returns(writer) 122 | expect(writer.write_bit).args('active') 123 | expect(mock(channel_class, 'MethodFrame')).args( 124 | 42, 20, 21, writer).returns('frame') 125 | expect(self.klass.send_frame).args('frame') 126 | expect(self.klass._flow_control_cb) 127 | 128 | self.klass._recv_flow(rframe) 129 | 130 | def test_recv_flow_ok_no_cb(self): 131 | self.klass._flow_control_cb = None 132 | rframe = mock() 133 | expect(rframe.args.read_bit).returns('active') 134 | 135 | self.klass._recv_flow_ok(rframe) 136 | assert_equals('active', self.klass.channel._active) 137 | 138 | def test_recv_flow_ok_with_cb(self): 139 | self.klass._flow_control_cb = mock() 140 | rframe = mock() 141 | expect(rframe.args.read_bit).returns('active') 142 | expect(self.klass._flow_control_cb) 143 | 144 | self.klass._recv_flow_ok(rframe) 145 | assert_equals('active', self.klass.channel._active) 146 | 147 | def test_close_when_not_closed(self): 148 | writer = mock() 149 | expect(mock(channel_class, 'Writer')).returns(writer) 150 | expect(writer.write_short).args('rcode') 151 | expect(writer.write_shortstr).args(('reason' * 60)[:255]) 152 | expect(writer.write_short).args('cid') 153 | expect(writer.write_short).args('mid') 154 | expect(mock(channel_class, 'MethodFrame')).args( 155 | 42, 20, 40, writer).returns('frame') 156 | expect(self.klass.send_frame).args('frame') 157 | expect(self.klass.channel.add_synchronous_cb).args( 158 | self.klass._recv_close_ok) 159 | 160 | self.klass.close('rcode', 'reason' * 60, 'cid', 'mid') 161 | assert_true(self.klass.channel._closed) 162 | assert_equals({ 163 | 'reply_code': 'rcode', 164 | 'reply_text': 'reason' * 60, 165 | 'class_id': 'cid', 166 | 'method_id': 'mid', 167 | }, self.klass.channel._close_info) 168 | 169 | def test_close_when_closed(self): 170 | self.klass.channel._closed = True 171 | stub(self.klass.send_frame) 172 | 173 | self.klass.close() 174 | 175 | def test_close_when_channel_reference_cleared_in_recv_close_ok(self): 176 | writer = mock() 177 | expect(mock(channel_class, 'Writer')).returns(writer) 178 | expect(writer.write_short).args('rcode') 179 | expect(writer.write_shortstr).args('reason') 180 | expect(writer.write_short).args('cid') 181 | expect(writer.write_short).args('mid') 182 | expect(mock(channel_class, 'MethodFrame')).args( 183 | 42, 20, 40, writer).returns('frame') 184 | expect(self.klass.send_frame).args('frame') 185 | expect(self.klass.channel.add_synchronous_cb).args(self.klass._recv_close_ok).side_effect( 186 | setattr, self.klass, '_channel', None) 187 | 188 | # assert nothing raised 189 | self.klass.close('rcode', 'reason', 'cid', 'mid') 190 | 191 | def test_close_when_error_sending_frame(self): 192 | self.klass.channel._closed = False 193 | writer = mock() 194 | expect(mock(channel_class, 'Writer')).returns(writer) 195 | expect(writer.write_short).args(0) 196 | expect(writer.write_shortstr).args('') 197 | expect(writer.write_short).args(0) 198 | expect(writer.write_short).args(0) 199 | expect(mock(channel_class, 'MethodFrame')).args( 200 | 42, 20, 40, writer).returns('frame') 201 | expect(self.klass.send_frame).args( 202 | 'frame').raises(RuntimeError('fail')) 203 | 204 | assert_raises(RuntimeError, self.klass.close) 205 | assert_true(self.klass.channel._closed) 206 | assert_equals({ 207 | 'reply_code': 0, 208 | 'reply_text': '', 209 | 'class_id': 0, 210 | 'method_id': 0, 211 | }, self.klass.channel._close_info) 212 | 213 | def test_recv_close(self): 214 | rframe = mock() 215 | expect(rframe.args.read_short).returns('rcode') 216 | expect(rframe.args.read_shortstr).returns('reason') 217 | expect(rframe.args.read_short).returns('cid') 218 | expect(rframe.args.read_short).returns('mid') 219 | 220 | expect(mock(channel_class, 'MethodFrame')).args( 221 | 42, 20, 41).returns('frame') 222 | expect(self.klass.channel._closed_cb).args(final_frame='frame') 223 | 224 | assert_false(self.klass.channel._closed) 225 | self.klass._recv_close(rframe) 226 | assert_true(self.klass.channel._closed) 227 | assert_equals({ 228 | 'reply_code': 'rcode', 229 | 'reply_text': 'reason', 230 | 'class_id': 'cid', 231 | 'method_id': 'mid', 232 | }, self.klass.channel._close_info) 233 | 234 | def test_recv_close_ok(self): 235 | expect(self.klass.channel._closed_cb) 236 | 237 | self.klass.channel._closed = False 238 | self.klass._recv_close_ok('frame') 239 | assert_true(self.klass.channel._closed) 240 | -------------------------------------------------------------------------------- /tests/unit/classes/exchange_class_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | from collections import deque 7 | 8 | from chai import Chai 9 | 10 | from haigha.classes import exchange_class 11 | from haigha.classes.protocol_class import ProtocolClass 12 | from haigha.classes.exchange_class import ExchangeClass 13 | from haigha.frames.method_frame import MethodFrame 14 | from haigha.writer import Writer 15 | 16 | 17 | class ExchangeClassTest(Chai): 18 | 19 | def setUp(self): 20 | super(ExchangeClassTest, self).setUp() 21 | ch = mock() 22 | ch.channel_id = 42 23 | ch.logger = mock() 24 | self.klass = ExchangeClass(ch) 25 | 26 | def test_init(self): 27 | expect(ProtocolClass.__init__).args('foo', a='b') 28 | 29 | klass = ExchangeClass.__new__(ExchangeClass) 30 | klass.__init__('foo', a='b') 31 | 32 | assert_equals( 33 | { 34 | 11: klass._recv_declare_ok, 35 | 21: klass._recv_delete_ok, 36 | }, klass.dispatch_map) 37 | assert_equals(deque(), klass._declare_cb) 38 | assert_equals(deque(), klass._delete_cb) 39 | 40 | def test_cleanup(self): 41 | self.klass._cleanup() 42 | assert_equals(None, self.klass._declare_cb) 43 | assert_equals(None, self.klass._delete_cb) 44 | assert_equals(None, self.klass._channel) 45 | assert_equals(None, self.klass.dispatch_map) 46 | 47 | def test_declare_default_args(self): 48 | w = mock() 49 | expect(self.klass.allow_nowait).returns(True) 50 | expect(mock(exchange_class, 'Writer')).returns(w) 51 | expect(w.write_short).args(self.klass.default_ticket).returns(w) 52 | expect(w.write_shortstr).args('exchange').returns(w) 53 | expect(w.write_shortstr).args('topic').returns(w) 54 | expect(w.write_bits).args(False, False, False, False, True).returns(w) 55 | expect(w.write_table).args({}) 56 | expect(mock(exchange_class, 'MethodFrame')).args( 57 | 42, 40, 10, w).returns('frame') 58 | expect(self.klass.send_frame).args('frame') 59 | stub(self.klass.channel.add_synchronous_cb) 60 | 61 | self.klass.declare('exchange', 'topic') 62 | assert_equals(deque(), self.klass._declare_cb) 63 | 64 | def test_declare_with_args(self): 65 | w = mock() 66 | stub(self.klass.allow_nowait) 67 | expect(mock(exchange_class, 'Writer')).returns(w) 68 | expect(w.write_short).args('t').returns(w) 69 | expect(w.write_shortstr).args('exchange').returns(w) 70 | expect(w.write_shortstr).args('topic').returns(w) 71 | expect(w.write_bits).args('p', 'd', False, False, False).returns(w) 72 | expect(w.write_table).args('table') 73 | expect(mock(exchange_class, 'MethodFrame')).args( 74 | 42, 40, 10, w).returns('frame') 75 | expect(self.klass.send_frame).args('frame') 76 | expect(self.klass.channel.add_synchronous_cb).args( 77 | self.klass._recv_declare_ok) 78 | 79 | self.klass.declare('exchange', 'topic', passive='p', durable='d', 80 | nowait=False, arguments='table', ticket='t') 81 | assert_equals(deque([None]), self.klass._declare_cb) 82 | 83 | def test_declare_with_cb(self): 84 | w = mock() 85 | expect(self.klass.allow_nowait).returns(True) 86 | expect(mock(exchange_class, 'Writer')).returns(w) 87 | expect(w.write_short).args('t').returns(w) 88 | expect(w.write_shortstr).args('exchange').returns(w) 89 | expect(w.write_shortstr).args('topic').returns(w) 90 | expect(w.write_bits).args('p', 'd', False, False, False).returns(w) 91 | expect(w.write_table).args('table') 92 | expect(mock(exchange_class, 'MethodFrame')).args( 93 | 42, 40, 10, w).returns('frame') 94 | expect(self.klass.send_frame).args('frame') 95 | expect(self.klass.channel.add_synchronous_cb).args( 96 | self.klass._recv_declare_ok) 97 | 98 | self.klass.declare('exchange', 'topic', passive='p', durable='d', 99 | nowait=True, arguments='table', ticket='t', cb='foo') 100 | assert_equals(deque(['foo']), self.klass._declare_cb) 101 | 102 | def test_recv_declare_ok_no_cb(self): 103 | self.klass._declare_cb = deque([None]) 104 | self.klass._recv_declare_ok('frame') 105 | assert_equals(deque(), self.klass._declare_cb) 106 | 107 | def test_recv_declare_ok_with_cb(self): 108 | cb = mock() 109 | self.klass._declare_cb = deque([cb]) 110 | expect(cb) 111 | self.klass._recv_declare_ok('frame') 112 | assert_equals(deque(), self.klass._declare_cb) 113 | 114 | def test_delete_default_args(self): 115 | w = mock() 116 | expect(self.klass.allow_nowait).returns(True) 117 | expect(mock(exchange_class, 'Writer')).returns(w) 118 | expect(w.write_short).args(self.klass.default_ticket).returns(w) 119 | expect(w.write_shortstr).args('exchange').returns(w) 120 | expect(w.write_bits).args(False, True) 121 | expect(mock(exchange_class, 'MethodFrame')).args( 122 | 42, 40, 20, w).returns('frame') 123 | expect(self.klass.send_frame).args('frame') 124 | stub(self.klass.channel.add_synchronous_cb) 125 | 126 | self.klass.delete('exchange') 127 | assert_equals(deque(), self.klass._delete_cb) 128 | 129 | def test_delete_with_args(self): 130 | w = mock() 131 | stub(self.klass.allow_nowait) 132 | expect(mock(exchange_class, 'Writer')).returns(w) 133 | expect(w.write_short).args('t').returns(w) 134 | expect(w.write_shortstr).args('exchange').returns(w) 135 | expect(w.write_bits).args('maybe', False) 136 | expect(mock(exchange_class, 'MethodFrame')).args( 137 | 42, 40, 20, w).returns('frame') 138 | expect(self.klass.send_frame).args('frame') 139 | expect(self.klass.channel.add_synchronous_cb).args( 140 | self.klass._recv_delete_ok) 141 | 142 | self.klass.delete( 143 | 'exchange', if_unused='maybe', nowait=False, ticket='t') 144 | assert_equals(deque([None]), self.klass._delete_cb) 145 | 146 | def test_delete_with_cb(self): 147 | w = mock() 148 | expect(self.klass.allow_nowait).returns(True) 149 | expect(mock(exchange_class, 'Writer')).returns(w) 150 | expect(w.write_short).args('t').returns(w) 151 | expect(w.write_shortstr).args('exchange').returns(w) 152 | expect(w.write_bits).args('maybe', False) 153 | expect(mock(exchange_class, 'MethodFrame')).args( 154 | 42, 40, 20, w).returns('frame') 155 | expect(self.klass.send_frame).args('frame') 156 | expect(self.klass.channel.add_synchronous_cb).args( 157 | self.klass._recv_delete_ok) 158 | 159 | self.klass.delete( 160 | 'exchange', if_unused='maybe', nowait=True, ticket='t', cb='foo') 161 | assert_equals(deque(['foo']), self.klass._delete_cb) 162 | 163 | def test_recv_delete_ok_no_cb(self): 164 | self.klass._delete_cb = deque([None]) 165 | self.klass._recv_delete_ok('frame') 166 | assert_equals(deque(), self.klass._delete_cb) 167 | 168 | def test_recv_delete_ok_with_cb(self): 169 | cb = mock() 170 | self.klass._delete_cb = deque([cb]) 171 | expect(cb) 172 | self.klass._recv_delete_ok('frame') 173 | assert_equals(deque(), self.klass._delete_cb) 174 | -------------------------------------------------------------------------------- /tests/unit/classes/protocol_class_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.classes import protocol_class 10 | from haigha.classes.protocol_class import ProtocolClass 11 | 12 | 13 | class ProtocolClassTest(Chai): 14 | 15 | def test_dispatch_when_in_dispatch_map(self): 16 | ch = mock() 17 | frame = mock() 18 | frame.method_id = 42 19 | 20 | klass = ProtocolClass(ch) 21 | klass.dispatch_map = {42: 'method'} 22 | 23 | with expect(ch.clear_synchronous_cb).args('method').returns(mock()) as cb: 24 | expect(cb).args(frame) 25 | 26 | klass.dispatch(frame) 27 | -------------------------------------------------------------------------------- /tests/unit/classes/transaction_class_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.classes import transaction_class 10 | from haigha.classes.protocol_class import ProtocolClass 11 | from haigha.classes.transaction_class import TransactionClass 12 | from haigha.frames.method_frame import MethodFrame 13 | from haigha.writer import Writer 14 | 15 | from collections import deque 16 | 17 | 18 | class TransactionClassTest(Chai): 19 | 20 | def setUp(self): 21 | super(TransactionClassTest, self).setUp() 22 | ch = mock() 23 | ch.channel_id = 42 24 | ch.logger = mock() 25 | self.klass = TransactionClass(ch) 26 | 27 | def test_init(self): 28 | expect(ProtocolClass.__init__).args('foo', a='b') 29 | 30 | klass = TransactionClass.__new__(TransactionClass) 31 | klass.__init__('foo', a='b') 32 | 33 | assert_equals( 34 | { 35 | 11: klass._recv_select_ok, 36 | 21: klass._recv_commit_ok, 37 | 31: klass._recv_rollback_ok, 38 | }, klass.dispatch_map) 39 | assert_false(klass._enabled) 40 | assert_equals(deque(), klass._select_cb) 41 | assert_equals(deque(), klass._commit_cb) 42 | assert_equals(deque(), klass._rollback_cb) 43 | 44 | def test_cleanup(self): 45 | self.klass._cleanup() 46 | assert_equals(None, self.klass._select_cb) 47 | assert_equals(None, self.klass._commit_cb) 48 | assert_equals(None, self.klass._rollback_cb) 49 | assert_equals(None, self.klass._channel) 50 | assert_equals(None, self.klass.dispatch_map) 51 | 52 | def test_properties(self): 53 | self.klass._enabled = 'maybe' 54 | assert_equals('maybe', self.klass.enabled) 55 | 56 | def test_select_when_not_enabled_and_no_cb(self): 57 | self.klass._enabled = False 58 | expect(mock(transaction_class, 'MethodFrame')).args( 59 | 42, 90, 10).returns('frame') 60 | expect(self.klass.send_frame).args('frame') 61 | expect(self.klass.channel.add_synchronous_cb).args( 62 | self.klass._recv_select_ok) 63 | 64 | self.klass.select() 65 | assert_true(self.klass.enabled) 66 | assert_equals(deque([None]), self.klass._select_cb) 67 | 68 | def test_select_when_not_enabled_with_cb(self): 69 | self.klass._enabled = False 70 | expect(mock(transaction_class, 'MethodFrame')).args( 71 | 42, 90, 10).returns('frame') 72 | expect(self.klass.send_frame).args('frame') 73 | expect(self.klass.channel.add_synchronous_cb).args( 74 | self.klass._recv_select_ok) 75 | 76 | self.klass.select(cb='foo') 77 | assert_true(self.klass.enabled) 78 | assert_equals(deque(['foo']), self.klass._select_cb) 79 | 80 | def test_select_when_already_enabled(self): 81 | self.klass._enabled = True 82 | stub(self.klass.send_frame) 83 | 84 | assert_equals(deque(), self.klass._select_cb) 85 | self.klass.select() 86 | assert_equals(deque(), self.klass._select_cb) 87 | 88 | def test_recv_select_ok_with_cb(self): 89 | cb = mock() 90 | self.klass._select_cb.append(cb) 91 | self.klass._select_cb.append(mock()) 92 | expect(cb) 93 | self.klass._recv_select_ok('frame') 94 | assert_equals(1, len(self.klass._select_cb)) 95 | assert_false(cb in self.klass._select_cb) 96 | 97 | def test_recv_select_ok_without_cb(self): 98 | self.klass._select_cb.append(None) 99 | self.klass._select_cb.append(mock()) 100 | 101 | self.klass._recv_select_ok('frame') 102 | assert_equals(1, len(self.klass._select_cb)) 103 | assert_false(None in self.klass._select_cb) 104 | 105 | def test_commit_when_enabled_no_cb(self): 106 | self.klass._enabled = True 107 | 108 | expect(mock(transaction_class, 'MethodFrame')).args( 109 | 42, 90, 20).returns('frame') 110 | expect(self.klass.send_frame).args('frame') 111 | expect(self.klass.channel.add_synchronous_cb).args( 112 | self.klass._recv_commit_ok) 113 | 114 | assert_equals(deque(), self.klass._commit_cb) 115 | self.klass.commit() 116 | assert_equals(deque([None]), self.klass._commit_cb) 117 | 118 | def test_commit_when_enabled_with_cb(self): 119 | self.klass._enabled = True 120 | 121 | expect(mock(transaction_class, 'MethodFrame')).args( 122 | 42, 90, 20).returns('frame') 123 | expect(self.klass.send_frame).args('frame') 124 | expect(self.klass.channel.add_synchronous_cb).args( 125 | self.klass._recv_commit_ok) 126 | 127 | self.klass._commit_cb = deque(['blargh']) 128 | self.klass.commit(cb='callback') 129 | assert_equals(deque(['blargh', 'callback']), self.klass._commit_cb) 130 | 131 | def test_commit_raises_transactionsnotenabled_when_not_enabled(self): 132 | self.klass._enabled = False 133 | assert_raises( 134 | TransactionClass.TransactionsNotEnabled, self.klass.commit) 135 | 136 | def test_recv_commit_ok_with_cb(self): 137 | cb = mock() 138 | self.klass._commit_cb.append(cb) 139 | self.klass._commit_cb.append(mock()) 140 | expect(cb) 141 | 142 | self.klass._recv_commit_ok('frame') 143 | assert_equals(1, len(self.klass._commit_cb)) 144 | assert_false(cb in self.klass._commit_cb) 145 | 146 | def test_recv_commit_ok_without_cb(self): 147 | self.klass._commit_cb.append(None) 148 | self.klass._commit_cb.append(mock()) 149 | 150 | self.klass._recv_commit_ok('frame') 151 | assert_equals(1, len(self.klass._commit_cb)) 152 | assert_false(None in self.klass._commit_cb) 153 | 154 | def test_rollback_when_enabled_no_cb(self): 155 | self.klass._enabled = True 156 | 157 | expect(mock(transaction_class, 'MethodFrame')).args( 158 | 42, 90, 30).returns('frame') 159 | expect(self.klass.send_frame).args('frame') 160 | expect(self.klass.channel.add_synchronous_cb).args( 161 | self.klass._recv_rollback_ok) 162 | 163 | assert_equals(deque(), self.klass._rollback_cb) 164 | self.klass.rollback() 165 | assert_equals(deque([None]), self.klass._rollback_cb) 166 | 167 | def test_rollback_when_enabled_with_cb(self): 168 | self.klass._enabled = True 169 | 170 | expect(mock(transaction_class, 'MethodFrame')).args( 171 | 42, 90, 30).returns('frame') 172 | expect(self.klass.send_frame).args('frame') 173 | expect(self.klass.channel.add_synchronous_cb).args( 174 | self.klass._recv_rollback_ok) 175 | 176 | self.klass._rollback_cb = deque(['blargh']) 177 | self.klass.rollback(cb='callback') 178 | assert_equals(deque(['blargh', 'callback']), self.klass._rollback_cb) 179 | 180 | def test_rollback_raises_transactionsnotenabled_when_not_enabled(self): 181 | self.klass._enabled = False 182 | assert_raises( 183 | TransactionClass.TransactionsNotEnabled, self.klass.rollback) 184 | 185 | def test_recv_rollback_ok_with_cb(self): 186 | cb = mock() 187 | self.klass._rollback_cb.append(cb) 188 | self.klass._rollback_cb.append(mock()) 189 | expect(cb) 190 | 191 | self.klass._recv_rollback_ok('frame') 192 | assert_equals(1, len(self.klass._rollback_cb)) 193 | assert_false(cb in self.klass._rollback_cb) 194 | 195 | def test_recv_rollback_ok_without_cb(self): 196 | self.klass._rollback_cb.append(None) 197 | self.klass._rollback_cb.append(mock()) 198 | 199 | self.klass._recv_rollback_ok('frame') 200 | assert_equals(1, len(self.klass._rollback_cb)) 201 | assert_false(None in self.klass._rollback_cb) 202 | -------------------------------------------------------------------------------- /tests/unit/frames/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/tests/unit/frames/__init__.py -------------------------------------------------------------------------------- /tests/unit/frames/content_frame_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.frames import content_frame 10 | from haigha.frames.content_frame import ContentFrame 11 | from haigha.frames.frame import Frame 12 | 13 | 14 | class ContentFrameTest(Chai): 15 | 16 | def test_type(self): 17 | assert_equals(3, ContentFrame.type()) 18 | 19 | def test_payload(self): 20 | klass = ContentFrame(42, 'payload') 21 | assert_equals('payload', klass.payload) 22 | 23 | def test_parse(self): 24 | frame = ContentFrame.parse(42, 'payload') 25 | assert_true(isinstance(frame, ContentFrame)) 26 | assert_equals(42, frame.channel_id) 27 | assert_equals('payload', frame.payload) 28 | 29 | def test_create_frames(self): 30 | itr = ContentFrame.create_frames(42, 'helloworld', 13) 31 | 32 | frame = itr.next() 33 | assert_true(isinstance(frame, ContentFrame)) 34 | assert_equals(42, frame.channel_id) 35 | assert_equals('hello', frame.payload) 36 | 37 | frame = itr.next() 38 | assert_true(isinstance(frame, ContentFrame)) 39 | assert_equals(42, frame.channel_id) 40 | assert_equals('world', frame.payload) 41 | 42 | assert_raises(StopIteration, itr.next) 43 | 44 | def test_init(self): 45 | expect(Frame.__init__).args(is_a(ContentFrame), 42) 46 | frame = ContentFrame(42, 'payload') 47 | assert_equals('payload', frame._payload) 48 | 49 | def test_str(self): 50 | # Test both branches but don't assert the actual content because its 51 | # not worth it 52 | frame = ContentFrame(42, 'payload') 53 | str(frame) 54 | 55 | frame = ContentFrame(42, 8675309) 56 | str(frame) 57 | 58 | def test_write_frame(self): 59 | w = mock() 60 | expect(mock(content_frame, 'Writer')).args('buffer').returns(w) 61 | expect(w.write_octet).args(3).returns(w) 62 | expect(w.write_short).args(42).returns(w) 63 | expect(w.write_long).args(5).returns(w) 64 | expect(w.write).args('hello').returns(w) 65 | expect(w.write_octet).args(0xce) 66 | 67 | frame = ContentFrame(42, 'hello') 68 | frame.write_frame('buffer') 69 | -------------------------------------------------------------------------------- /tests/unit/frames/frame_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | import struct 9 | from collections import deque 10 | 11 | from haigha.frames import frame 12 | from haigha.frames.frame import Frame 13 | from haigha.reader import Reader 14 | 15 | 16 | class FrameTest(Chai): 17 | 18 | def test_register(self): 19 | class DummyFrame(Frame): 20 | 21 | @classmethod 22 | def type(self): 23 | return 42 24 | 25 | assertEquals(None, Frame._frame_type_map.get(42)) 26 | DummyFrame.register() 27 | assertEquals(DummyFrame, Frame._frame_type_map[42]) 28 | 29 | def test_type_raises_not_implemented(self): 30 | assertRaises(NotImplementedError, Frame.type) 31 | 32 | def test_read_frames_reads_until_buffer_underflow(self): 33 | reader = mock() 34 | 35 | expect(reader.tell).returns(0) 36 | expect(Frame._read_frame).args(reader).returns('frame1') 37 | 38 | expect(reader.tell).returns(2) 39 | expect(Frame._read_frame).args(reader).returns('frame2') 40 | 41 | expect(reader.tell).returns(3) 42 | expect(Frame._read_frame).args(reader).raises(Reader.BufferUnderflow) 43 | 44 | expect(reader.seek).args(3) 45 | 46 | assertEquals(deque(['frame1', 'frame2']), Frame.read_frames(reader)) 47 | 48 | def test_read_frames_handles_reader_errors(self): 49 | reader = mock() 50 | self.mock(Frame, '_read_frame') 51 | 52 | expect(reader.tell).returns(0) 53 | expect(Frame._read_frame).args( 54 | reader).raises(Reader.ReaderError("bad!")) 55 | 56 | assertRaises(Frame.FormatError, Frame.read_frames, reader) 57 | 58 | def test_read_frames_handles_struct_errors(self): 59 | reader = mock() 60 | self.mock(Frame, '_read_frame') 61 | 62 | expect(reader.tell).returns(0) 63 | expect(Frame._read_frame).args(reader).raises(struct.error("bad!")) 64 | 65 | self.assertRaises(Frame.FormatError, Frame.read_frames, reader) 66 | 67 | def test_read_frame_on_full_frame(self): 68 | class FrameReader(Frame): 69 | 70 | @classmethod 71 | def type(self): 72 | return 45 73 | 74 | @classmethod 75 | def parse(self, channel_id, payload): 76 | return 'no_frame' 77 | FrameReader.register() 78 | 79 | self.mock(frame, 'Reader') 80 | reader = self.mock() 81 | payload = self.mock() 82 | 83 | expect(reader.read_octet).returns(45) # frame type 84 | expect(reader.read_short).returns(32) # channel id 85 | expect(reader.read_long).returns(42) # size 86 | 87 | expect(reader.tell).returns(5) 88 | expect(frame.Reader).args(reader, 5, 42).returns(payload) 89 | expect(reader.seek).args(42, 1) 90 | 91 | expect(reader.read_octet).returns(0xce) 92 | expect(FrameReader.parse).args(32, payload).returns('a_frame') 93 | 94 | assertEquals('a_frame', Frame._read_frame(reader)) 95 | 96 | def test_read_frame_raises_bufferunderflow_when_incomplete_payload(self): 97 | self.mock(frame, 'Reader') 98 | reader = self.mock() 99 | 100 | expect(reader.read_octet).returns(45) # frame type 101 | expect(reader.read_short).returns(32) # channel id 102 | expect(reader.read_long).returns(42) # size 103 | 104 | expect(reader.tell).returns(5) 105 | expect(frame.Reader).args(reader, 5, 42).returns('payload') 106 | expect(reader.seek).args(42, 1) 107 | 108 | expect(reader.read_octet).raises(Reader.BufferUnderflow) 109 | assert_raises(Reader.BufferUnderflow, Frame._read_frame, reader) 110 | 111 | def test_read_frame_raises_formaterror_if_bad_footer(self): 112 | self.mock(frame, 'Reader') 113 | reader = self.mock() 114 | 115 | expect(reader.read_octet).returns(45) # frame type 116 | expect(reader.read_short).returns(32) # channel id 117 | expect(reader.read_long).returns(42) # size 118 | 119 | expect(reader.tell).returns(5) 120 | expect(frame.Reader).args(reader, 5, 42).returns('payload') 121 | expect(reader.seek).args(42, 1) 122 | expect(reader.read_octet).returns(0xff) 123 | 124 | assert_raises(Frame.FormatError, Frame._read_frame, reader) 125 | 126 | def test_read_frame_raises_invalidframetype_for_unregistered_frame_type(self): 127 | self.mock(frame, 'Reader') 128 | reader = self.mock() 129 | payload = self.mock() 130 | 131 | expect(reader.read_octet).returns(54) # frame type 132 | expect(reader.read_short).returns(32) # channel id 133 | expect(reader.read_long).returns(42) # size 134 | 135 | expect(reader.tell).returns(5) 136 | expect(frame.Reader).args(reader, 5, 42).returns(payload) 137 | expect(reader.seek).args(42, 1) 138 | 139 | expect(reader.read_octet).returns(0xce) 140 | 141 | assertRaises(Frame.InvalidFrameType, Frame._read_frame, reader) 142 | 143 | def test_parse_raises_not_implemented(self): 144 | assertRaises(NotImplementedError, Frame.parse, 'channel_id', 'payload') 145 | 146 | def test_properties(self): 147 | frame = Frame('channel_id') 148 | assert_equals('channel_id', frame.channel_id) 149 | 150 | def test_str(self): 151 | frame = Frame(42) 152 | assert_equals('Frame[channel: 42]', str(frame)) 153 | 154 | def test_repr(self): 155 | expect(Frame.__str__).returns('foo') 156 | frame = Frame(42) 157 | assert_equals('foo', repr(frame)) 158 | 159 | def test_write_frame(self): 160 | frame = Frame(42) 161 | assert_raises(NotImplementedError, frame.write_frame, 'stream') 162 | -------------------------------------------------------------------------------- /tests/unit/frames/header_frame_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | import struct 9 | import time 10 | from datetime import datetime 11 | 12 | from haigha.frames import header_frame 13 | from haigha.frames.header_frame import HeaderFrame 14 | from haigha.reader import Reader 15 | from haigha.writer import Writer 16 | 17 | 18 | class HeaderFrameTest(Chai): 19 | 20 | def test_type(self): 21 | assert_equals(2, HeaderFrame.type()) 22 | 23 | def test_properties(self): 24 | frame = HeaderFrame(42, 'class_id', 'weight', 'size', 'props') 25 | assert_equals(42, frame.channel_id) 26 | assert_equals('class_id', frame.class_id) 27 | assert_equals('weight', frame.weight) 28 | assert_equals('size', frame.size) 29 | assert_equals('props', frame.properties) 30 | 31 | def test_str(self): 32 | # Don't bother checking the copy 33 | frame = HeaderFrame(42, 5, 6, 7, 'props') 34 | assert_equals('HeaderFrame[channel: 42, class_id: 5, weight: 6, size: 7, properties: props]', 35 | str(frame)) 36 | 37 | def test_parse_fast_for_standard_properties(self): 38 | bit_writer = Writer() 39 | val_writer = Writer() 40 | 41 | # strip ms because amqp doesn't include it 42 | now = datetime.utcfromtimestamp( 43 | long(time.mktime(datetime.now().timetuple()))) 44 | 45 | bit_field = 0 46 | for pname, ptype, reader, writer, mask in HeaderFrame.PROPERTIES: 47 | bit_field = (bit_field << 1) | 1 48 | 49 | if ptype == 'shortstr': 50 | val_writer.write_shortstr(pname) 51 | elif ptype == 'octet': 52 | val_writer.write_octet(42) 53 | elif ptype == 'timestamp': 54 | val_writer.write_timestamp(now) 55 | elif ptype == 'table': 56 | val_writer.write_table({'foo': 'bar'}) 57 | 58 | bit_field <<= (16 - len(HeaderFrame.PROPERTIES)) 59 | bit_writer.write_short(bit_field) 60 | 61 | header_writer = Writer() 62 | header_writer.write_short(5) 63 | header_writer.write_short(6) 64 | header_writer.write_longlong(7) 65 | payload = header_writer.buffer() 66 | payload += bit_writer.buffer() 67 | payload += val_writer.buffer() 68 | 69 | reader = Reader(payload) 70 | frame = HeaderFrame.parse(4, reader) 71 | 72 | for pname, ptype, reader, writer, mask in HeaderFrame.PROPERTIES: 73 | if ptype == 'shortstr': 74 | self.assertEquals(pname, frame.properties[pname]) 75 | elif ptype == 'octet': 76 | self.assertEquals(42, frame.properties[pname]) 77 | elif ptype == 'timestamp': 78 | self.assertEquals(now, frame.properties[pname]) 79 | elif ptype == 'table': 80 | self.assertEquals({'foo': 'bar'}, frame.properties[pname]) 81 | 82 | assert_equals(4, frame.channel_id) 83 | assert_equals(5, frame._class_id) 84 | assert_equals(6, frame._weight) 85 | assert_equals(7, frame._size) 86 | 87 | def test_parse_slow_for_standard_properties(self): 88 | HeaderFrame.DEFAULT_PROPERTIES = False 89 | bit_writer = Writer() 90 | val_writer = Writer() 91 | 92 | # strip ms because amqp doesn't include it 93 | now = datetime.utcfromtimestamp( 94 | long(time.mktime(datetime.now().timetuple()))) 95 | 96 | bit_field = 0 97 | for pname, ptype, reader, writer, mask in HeaderFrame.PROPERTIES: 98 | bit_field = (bit_field << 1) | 1 99 | 100 | if ptype == 'shortstr': 101 | val_writer.write_shortstr(pname) 102 | elif ptype == 'octet': 103 | val_writer.write_octet(42) 104 | elif ptype == 'timestamp': 105 | val_writer.write_timestamp(now) 106 | elif ptype == 'table': 107 | val_writer.write_table({'foo': 'bar'}) 108 | 109 | bit_field <<= (16 - len(HeaderFrame.PROPERTIES)) 110 | bit_writer.write_short(bit_field) 111 | 112 | header_writer = Writer() 113 | header_writer.write_short(5) 114 | header_writer.write_short(6) 115 | header_writer.write_longlong(7) 116 | payload = header_writer.buffer() 117 | payload += bit_writer.buffer() 118 | payload += val_writer.buffer() 119 | 120 | reader = Reader(payload) 121 | frame = HeaderFrame.parse(4, reader) 122 | HeaderFrame.DEFAULT_PROPERTIES = True 123 | 124 | for pname, ptype, reader, writer, mask in HeaderFrame.PROPERTIES: 125 | if ptype == 'shortstr': 126 | self.assertEquals(pname, frame.properties[pname]) 127 | elif ptype == 'octet': 128 | self.assertEquals(42, frame.properties[pname]) 129 | elif ptype == 'timestamp': 130 | self.assertEquals(now, frame.properties[pname]) 131 | elif ptype == 'table': 132 | self.assertEquals({'foo': 'bar'}, frame.properties[pname]) 133 | 134 | assert_equals(4, frame.channel_id) 135 | assert_equals(5, frame._class_id) 136 | assert_equals(6, frame._weight) 137 | assert_equals(7, frame._size) 138 | 139 | def test_write_frame_fast_for_standard_properties(self): 140 | bit_field = 0 141 | properties = {} 142 | now = datetime.utcfromtimestamp( 143 | long(time.mktime(datetime.now().timetuple()))) 144 | for pname, ptype, reader, writer, mask in HeaderFrame.PROPERTIES: 145 | bit_field |= mask 146 | 147 | if ptype == 'shortstr': 148 | properties[pname] = pname 149 | elif ptype == 'octet': 150 | properties[pname] = 42 151 | elif ptype == 'timestamp': 152 | properties[pname] = now 153 | elif ptype == 'table': 154 | properties[pname] = {'foo': 'bar'} 155 | 156 | frame = HeaderFrame(42, 5, 6, 7, properties) 157 | buf = bytearray() 158 | frame.write_frame(buf) 159 | 160 | reader = Reader(buf) 161 | assert_equals(2, reader.read_octet()) 162 | assert_equals(42, reader.read_short()) 163 | size = reader.read_long() 164 | start_pos = reader.tell() 165 | assert_equals(5, reader.read_short()) 166 | assert_equals(6, reader.read_short()) 167 | assert_equals(7, reader.read_longlong()) 168 | assert_equals(0b1111111111111100, reader.read_short()) 169 | 170 | for pname, ptype, rfunc, wfunc, mask in HeaderFrame.PROPERTIES: 171 | if ptype == 'shortstr': 172 | assertEquals(pname, reader.read_shortstr()) 173 | elif ptype == 'octet': 174 | assertEquals(42, reader.read_octet()) 175 | elif ptype == 'timestamp': 176 | assertEquals(now, reader.read_timestamp()) 177 | elif ptype == 'table': 178 | assertEquals({'foo': 'bar'}, reader.read_table()) 179 | 180 | end_pos = reader.tell() 181 | assert_equals(size, end_pos - start_pos) 182 | assert_equals(0xce, reader.read_octet()) 183 | 184 | def test_write_frame_slow_for_standard_properties(self): 185 | HeaderFrame.DEFAULT_PROPERTIES = False 186 | bit_field = 0 187 | properties = {} 188 | now = datetime.utcfromtimestamp( 189 | long(time.mktime(datetime.now().timetuple()))) 190 | for pname, ptype, reader, writer, mask in HeaderFrame.PROPERTIES: 191 | bit_field |= mask 192 | 193 | if ptype == 'shortstr': 194 | properties[pname] = pname 195 | elif ptype == 'octet': 196 | properties[pname] = 42 197 | elif ptype == 'timestamp': 198 | properties[pname] = now 199 | elif ptype == 'table': 200 | properties[pname] = {'foo': 'bar'} 201 | 202 | frame = HeaderFrame(42, 5, 6, 7, properties) 203 | buf = bytearray() 204 | frame.write_frame(buf) 205 | HeaderFrame.DEFAULT_PROPERTIES = True 206 | 207 | reader = Reader(buf) 208 | assert_equals(2, reader.read_octet()) 209 | assert_equals(42, reader.read_short()) 210 | size = reader.read_long() 211 | start_pos = reader.tell() 212 | assert_equals(5, reader.read_short()) 213 | assert_equals(6, reader.read_short()) 214 | assert_equals(7, reader.read_longlong()) 215 | assert_equals(0b1111111111111100, reader.read_short()) 216 | 217 | for pname, ptype, rfunc, wfunc, mask in HeaderFrame.PROPERTIES: 218 | if ptype == 'shortstr': 219 | assertEquals(pname, reader.read_shortstr()) 220 | elif ptype == 'octet': 221 | assertEquals(42, reader.read_octet()) 222 | elif ptype == 'timestamp': 223 | assertEquals(now, reader.read_timestamp()) 224 | elif ptype == 'table': 225 | assertEquals({'foo': 'bar'}, reader.read_table()) 226 | 227 | end_pos = reader.tell() 228 | assert_equals(size, end_pos - start_pos) 229 | assert_equals(0xce, reader.read_octet()) 230 | -------------------------------------------------------------------------------- /tests/unit/frames/heartbeat_frame_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.frames import heartbeat_frame 10 | from haigha.frames.heartbeat_frame import HeartbeatFrame 11 | from haigha.frames.frame import Frame 12 | 13 | 14 | class HeartbeatFrameTest(Chai): 15 | 16 | def test_type(self): 17 | assert_equals(8, HeartbeatFrame.type()) 18 | 19 | def test_parse(self): 20 | frame = HeartbeatFrame.parse(42, 'payload') 21 | assert_true(isinstance(frame, HeartbeatFrame)) 22 | assert_equals(42, frame.channel_id) 23 | 24 | def test_write_frame(self): 25 | w = mock() 26 | expect(mock(heartbeat_frame, 'Writer')).args('buffer').returns(w) 27 | expect(w.write_octet).args(8).returns(w) 28 | expect(w.write_short).args(42).returns(w) 29 | expect(w.write_long).args(0).returns(w) 30 | expect(w.write_octet).args(0xce) 31 | 32 | frame = HeartbeatFrame(42) 33 | frame.write_frame('buffer') 34 | -------------------------------------------------------------------------------- /tests/unit/frames/method_frame_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | import struct 9 | import time 10 | from datetime import datetime 11 | 12 | from haigha.frames import method_frame 13 | from haigha.frames.method_frame import MethodFrame 14 | from haigha.reader import Reader 15 | from haigha.writer import Writer 16 | 17 | 18 | class MethodFrameTest(Chai): 19 | 20 | def test_type(self): 21 | assert_equals(1, MethodFrame.type()) 22 | 23 | def test_properties(self): 24 | frame = MethodFrame('channel_id', 'class_id', 'method_id', 'args') 25 | assert_equals('channel_id', frame.channel_id) 26 | assert_equals('class_id', frame.class_id) 27 | assert_equals('method_id', frame.method_id) 28 | assert_equals('args', frame.args) 29 | 30 | def test_parse(self): 31 | reader = mock() 32 | expect(reader.read_short).returns('class_id') 33 | expect(reader.read_short).returns('method_id') 34 | frame = MethodFrame.parse(42, reader) 35 | 36 | assert_equals(42, frame.channel_id) 37 | assert_equals('class_id', frame.class_id) 38 | assert_equals('method_id', frame.method_id) 39 | assert_equals(reader, frame.args) 40 | 41 | def test_str(self): 42 | frame = MethodFrame(42, 5, 6, Reader(bytearray('hello'))) 43 | assert_equals( 44 | 'MethodFrame[channel: 42, class_id: 5, method_id: 6, args: \\x68\\x65\\x6c\\x6c\\x6f]', str(frame)) 45 | 46 | frame = MethodFrame(42, 5, 6) 47 | assert_equals( 48 | 'MethodFrame[channel: 42, class_id: 5, method_id: 6, args: None]', str(frame)) 49 | 50 | def test_write_frame(self): 51 | args = mock() 52 | expect(args.buffer).returns('hello') 53 | 54 | frame = MethodFrame(42, 5, 6, args) 55 | buf = bytearray() 56 | frame.write_frame(buf) 57 | 58 | reader = Reader(buf) 59 | assert_equals(1, reader.read_octet()) 60 | assert_equals(42, reader.read_short()) 61 | size = reader.read_long() 62 | start_pos = reader.tell() 63 | assert_equals(5, reader.read_short()) 64 | assert_equals(6, reader.read_short()) 65 | args_pos = reader.tell() 66 | assert_equals('hello', reader.read(size - (args_pos - start_pos))) 67 | assert_equals(0xce, reader.read_octet()) 68 | -------------------------------------------------------------------------------- /tests/unit/message_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.message import Message 10 | 11 | 12 | class MessageTest(Chai): 13 | 14 | def test_init_no_args(self): 15 | m = Message() 16 | self.assertEquals('', m._body) 17 | self.assertEquals(None, m._delivery_info) 18 | self.assertEquals(None, m.return_info) 19 | self.assertEquals({}, m._properties) 20 | 21 | def test_init_with_delivery_and_args(self): 22 | m = Message('foo', 'delivery', foo='bar') 23 | self.assertEquals('foo', m._body) 24 | self.assertEquals('delivery', m._delivery_info) 25 | self.assertEquals({'foo': 'bar'}, m._properties) 26 | 27 | m = Message(u'D\xfcsseldorf') 28 | self.assertEquals('D\xc3\xbcsseldorf', m._body) 29 | self.assertEquals({'content_encoding': 'utf-8'}, m._properties) 30 | 31 | def test_with_body_and_properties(self): 32 | m = Message('foo', foo='bar') 33 | self.assertEquals('foo', m.body) 34 | self.assertEquals(None, m.delivery_info) 35 | self.assertEquals(None, m.return_info) 36 | self.assertEquals({'foo': 'bar'}, m.properties) 37 | 38 | def test_with_delivery_and_properties(self): 39 | m = Message('foo', 'delivery', foo='bar') 40 | self.assertEquals('foo', m.body) 41 | self.assertEquals('delivery', m.delivery_info) 42 | self.assertEquals(None, m.return_info) 43 | self.assertEquals({'foo': 'bar'}, m.properties) 44 | 45 | def test_with_return_and_properties(self): 46 | m = Message('foo', return_info='return', foo='bar') 47 | self.assertEquals('foo', m.body) 48 | self.assertEquals('return', m.return_info) 49 | self.assertEquals(None, m.delivery_info) 50 | self.assertEquals({'foo': 'bar'}, m.properties) 51 | 52 | def test_len(self): 53 | m = Message('foobar') 54 | self.assertEquals(6, len(m)) 55 | 56 | def test_nonzero(self): 57 | m = Message() 58 | self.assertTrue(m) 59 | 60 | def test_eq(self): 61 | l = Message() 62 | r = Message() 63 | self.assertEquals(l, r) 64 | 65 | l = Message('foo') 66 | r = Message('foo') 67 | self.assertEquals(l, r) 68 | 69 | l = Message(foo='bar') 70 | r = Message(foo='bar') 71 | self.assertEquals(l, r) 72 | 73 | l = Message('hello', foo='bar') 74 | r = Message('hello', foo='bar') 75 | self.assertEquals(l, r) 76 | 77 | l = Message('foo') 78 | r = Message('bar') 79 | self.assertNotEquals(l, r) 80 | 81 | l = Message(foo='bar') 82 | r = Message(foo='brah') 83 | self.assertNotEquals(l, r) 84 | 85 | l = Message('hello', foo='bar') 86 | r = Message('goodbye', foo='bar') 87 | self.assertNotEquals(l, r) 88 | 89 | l = Message('hello', foo='bar') 90 | r = Message('hello', foo='brah') 91 | self.assertNotEquals(l, r) 92 | 93 | self.assertNotEquals(Message(), object()) 94 | 95 | def test_str_with_delivery_info(self): 96 | m = Message('foo', 'delivery', foo='bar') 97 | str(m) 98 | 99 | def test_str_with_return_info(self): 100 | m = Message('foo', return_info='returned', foo='bar') 101 | str(m) 102 | -------------------------------------------------------------------------------- /tests/unit/transports/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/agoragames/haigha/7b004e1c0316ec14b94fec1c54554654c38b1a25/tests/unit/transports/__init__.py -------------------------------------------------------------------------------- /tests/unit/transports/event_transport_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.transports import event_transport 10 | from haigha.transports.event_transport import * 11 | 12 | 13 | class EventTransportTest(Chai): 14 | 15 | def setUp(self): 16 | super(EventTransportTest, self).setUp() 17 | 18 | self.connection = mock() 19 | self.transport = EventTransport(self.connection) 20 | self.transport._host = 'server' 21 | 22 | def test_sock_close_cb(self): 23 | expect(self.connection.transport_closed).args( 24 | msg='socket to server closed unexpectedly') 25 | self.transport._sock_close_cb('sock') 26 | 27 | def test_sock_error_cb(self): 28 | expect(self.connection.transport_closed).args( 29 | msg='error on connection to server: amsg') 30 | self.transport._sock_error_cb('sock', 'amsg') 31 | 32 | def test_sock_read_cb(self): 33 | expect(self.connection.read_frames) 34 | self.transport._sock_read_cb('sock') 35 | 36 | def test_connect(self): 37 | sock = mock() 38 | mock(event_transport, 'EventSocket') 39 | self.connection._connect_timeout = 4.12 40 | self.connection._sock_opts = { 41 | ('family', 'tcp'): 34, 42 | ('range', 'ipv6'): 'hex' 43 | } 44 | 45 | expect(event_transport.EventSocket).args( 46 | read_cb=self.transport._sock_read_cb, 47 | close_cb=self.transport._sock_close_cb, 48 | error_cb=self.transport._sock_error_cb, 49 | debug=self.connection.debug, 50 | logger=self.connection.logger, 51 | ).returns(sock) 52 | expect(sock.setsockopt).args('family', 'tcp', 34).any_order() 53 | expect(sock.setsockopt).args('range', 'ipv6', 'hex').any_order() 54 | expect(sock.setblocking).args(False) 55 | expect(sock.connect).args(('host', 5309), timeout=4.12) 56 | 57 | self.transport.connect(('host', 5309)) 58 | 59 | def test_read(self): 60 | self.transport._heartbeat_timeout = None 61 | self.transport._sock = mock() 62 | expect(self.transport._sock.read).returns('buffereddata') 63 | assert_equals('buffereddata', self.transport.read()) 64 | 65 | def test_read_with_timeout_and_no_current_one(self): 66 | self.transport._heartbeat_timeout = None 67 | self.transport._sock = mock() 68 | mock(event_transport, 'event') 69 | expect(event_transport.event.timeout).args( 70 | 'timeout', self.transport._sock_read_cb, self.transport._sock).returns( 71 | 'timer') 72 | 73 | expect(self.transport._sock.read).returns('buffereddata') 74 | assert_equals('buffereddata', self.transport.read('timeout')) 75 | assert_equals('timer', self.transport._heartbeat_timeout) 76 | 77 | def test_read_with_timeout_and_current_one(self): 78 | self.transport._heartbeat_timeout = mock() 79 | self.transport._sock = mock() 80 | mock(event_transport, 'event') 81 | expect(self.transport._heartbeat_timeout.delete) 82 | expect(event_transport.event.timeout).args( 83 | 'timeout', self.transport._sock_read_cb, self.transport._sock).returns( 84 | 'timer') 85 | 86 | expect(self.transport._sock.read).returns('buffereddata') 87 | assert_equals('buffereddata', self.transport.read('timeout')) 88 | assert_equals('timer', self.transport._heartbeat_timeout) 89 | 90 | def test_read_without_timeout_but_current_one(self): 91 | self.transport._heartbeat_timeout = mock() 92 | self.transport._sock = mock() 93 | mock(event_transport, 'event') 94 | expect(self.transport._heartbeat_timeout.delete) 95 | 96 | expect(self.transport._sock.read).returns('buffereddata') 97 | assert_equals('buffereddata', self.transport.read()) 98 | assert_equals(None, self.transport._heartbeat_timeout) 99 | 100 | def test_read_when_no_sock(self): 101 | self.transport.read() 102 | 103 | def test_buffer(self): 104 | self.transport._sock = mock() 105 | expect(self.transport._sock.buffer).args('somedata') 106 | self.transport.buffer('somedata') 107 | 108 | def test_buffer_when_no_sock(self): 109 | self.transport.buffer('somedata') 110 | 111 | def test_write(self): 112 | self.transport._sock = mock() 113 | expect(self.transport._sock.write).args('somedata') 114 | self.transport.write('somedata') 115 | 116 | def test_write_when_no_sock(self): 117 | self.transport.write('somedata') 118 | 119 | def test_disconnect(self): 120 | self.transport._sock = mock() 121 | self.transport._sock.close_cb = 'cb' 122 | expect(self.transport._sock.close) 123 | self.transport.disconnect() 124 | assert_equals(None, self.transport._sock.close_cb) 125 | 126 | def test_disconnect_when_no_sock(self): 127 | self.transport.disconnect() 128 | -------------------------------------------------------------------------------- /tests/unit/transports/gevent_transport_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | import errno 9 | import unittest 10 | 11 | try: 12 | import gevent 13 | from gevent.coros import Semaphore 14 | from gevent import socket 15 | from gevent.pool import Pool 16 | 17 | from haigha.transports import gevent_transport 18 | from haigha.transports.gevent_transport import * 19 | except ImportError: 20 | import warnings 21 | warnings.warn('Failed to load gevent modules') 22 | gevent = None 23 | 24 | @unittest.skipIf(gevent is None, 'skipping gevent tests') 25 | class GeventTransportTest(Chai): 26 | 27 | def setUp(self): 28 | super(GeventTransportTest, self).setUp() 29 | 30 | self.connection = mock() 31 | self.transport = GeventTransport(self.connection) 32 | self.transport._host = 'server:1234' 33 | 34 | def test_init(self): 35 | assert_equals(bytearray(), self.transport._buffer) 36 | assert_true(isinstance(self.transport._read_lock, Semaphore)) 37 | assert_true(isinstance(self.transport._write_lock, Semaphore)) 38 | 39 | def test_connect(self): 40 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 41 | expect(parent.connect).args( 42 | ('host', 'port'), klass=is_arg(socket.socket)).returns('somedata') 43 | 44 | self.transport.connect(('host', 'port')) 45 | 46 | def test_read(self): 47 | #self.transport._read_lock = mock() 48 | #expect( self.transport._read_lock.locked ).returns( False ) 49 | expect(self.transport._read_lock.acquire) 50 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 51 | expect(parent.read).args(timeout=None).returns('somedata') 52 | expect(self.transport._read_lock.release) 53 | expect(self.transport._read_wait.set) 54 | expect(self.transport._read_wait.clear) 55 | 56 | assert_equals('somedata', self.transport.read()) 57 | 58 | def test_read_when_already_locked(self): 59 | expect(self.transport._read_lock.locked).returns(True) 60 | stub(self.transport._read_lock.acquire) 61 | stub(mock(gevent_transport, 'super')) 62 | stub(self.transport._read_lock.release) 63 | expect(self.transport._read_wait.wait) 64 | 65 | assert_equals(None, self.transport.read()) 66 | 67 | def test_read_when_raises_exception(self): 68 | #self.transport._read_lock = mock() 69 | expect(self.transport._read_lock.acquire) 70 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 71 | expect(parent.read).args(timeout='5').raises(Exception('fail')) 72 | expect(self.transport._read_lock.release) 73 | 74 | assert_raises(Exception, self.transport.read, timeout='5') 75 | 76 | def test_buffer(self): 77 | #self.transport._read_lock = mock() 78 | expect(self.transport._read_lock.acquire) 79 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 80 | expect(parent.buffer).args('datas') 81 | expect(self.transport._read_lock.release) 82 | 83 | self.transport.buffer('datas') 84 | 85 | def test_buffer_when_raises_exception(self): 86 | #self.transport._read_lock = mock() 87 | expect(self.transport._read_lock.acquire) 88 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 89 | expect(parent.buffer).args('datas').raises(Exception('fail')) 90 | expect(self.transport._read_lock.release) 91 | 92 | assert_raises(Exception, self.transport.buffer, 'datas') 93 | 94 | def test_write(self): 95 | #self.transport._write_lock = mock() 96 | expect(self.transport._write_lock.acquire) 97 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 98 | expect(parent.write).args('datas') 99 | expect(self.transport._write_lock.release) 100 | 101 | self.transport.write('datas') 102 | 103 | def test_write_when_raises_an_exception(self): 104 | #self.transport._write_lock = mock() 105 | expect(self.transport._write_lock.acquire) 106 | with expect(mock(gevent_transport, 'super')).args(is_arg(GeventTransport), GeventTransport).returns(mock()) as parent: 107 | expect(parent.write).args('datas').raises(Exception('fail')) 108 | expect(self.transport._write_lock.release) 109 | 110 | assert_raises(Exception, self.transport.write, 'datas') 111 | 112 | @unittest.skipIf(gevent is None, 'skipping gevent tests') 113 | class GeventPoolTransportTest(Chai): 114 | 115 | def setUp(self): 116 | super(GeventPoolTransportTest, self).setUp() 117 | 118 | self.connection = mock() 119 | self.transport = GeventPoolTransport(self.connection) 120 | self.transport._host = 'server:1234' 121 | 122 | def test_init(self): 123 | assert_equals(bytearray(), self.transport._buffer) 124 | assert_true(isinstance(self.transport._read_lock, Semaphore)) 125 | assert_true(isinstance(self.transport.pool, Pool)) 126 | 127 | trans = GeventPoolTransport(self.connection, pool='inground') 128 | assert_equals('inground', trans.pool) 129 | 130 | def test_process_channels(self): 131 | chs = [mock(), mock()] 132 | self.transport._pool = mock() 133 | 134 | expect(self.transport._pool.spawn).args(chs[0].process_frames) 135 | expect(self.transport._pool.spawn).args(chs[1].process_frames) 136 | 137 | self.transport.process_channels(chs) 138 | -------------------------------------------------------------------------------- /tests/unit/transports/transport_test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (c) 2011-2017, Agora Games, LLC All rights reserved. 3 | 4 | https://github.com/agoragames/haigha/blob/master/LICENSE.txt 5 | ''' 6 | 7 | from chai import Chai 8 | 9 | from haigha.transports.transport import Transport 10 | 11 | 12 | class TransportTest(Chai): 13 | 14 | def test_init_and_connection_property(self): 15 | t = Transport('conn') 16 | assert_equals('conn', t._connection) 17 | assert_equals('conn', t.connection) 18 | 19 | def test_process_channels(self): 20 | t = Transport('conn') 21 | ch1 = mock() 22 | ch2 = mock() 23 | chs = set([ch1, ch2]) 24 | expect(ch1.process_frames) 25 | expect(ch2.process_frames) 26 | 27 | t.process_channels(chs) 28 | --------------------------------------------------------------------------------