define claim/cancel message types

Author: blampe

tchannel/messages/__init__.py

@@ -20,28 +20,32 @@
 
 from __future__ import absolute_import
 
-from .types import Types
-from .error import ErrorMessage, ErrorCode, error_rw
-from .common import Tracing, ChecksumType
 from .call_request import CallRequestMessage, call_req_rw
+from .call_request_continue import call_req_c_rw
 from .call_response import CallResponseMessage, call_res_rw
+from .call_response_continue import call_res_c_rw
+from .cancel import CancelMessage, cancel_rw
+from .claim import ClaimMessage, claim_rw
+from .common import Tracing, ChecksumType
+from .error import ErrorMessage, ErrorCode, error_rw
 from .init_request import InitRequestMessage, init_req_rw
 from .init_response import InitResponseMessage, init_res_rw
 from .ping_request import PingRequestMessage, ping_req_rw
 from .ping_response import PingResponseMessage, ping_res_rw
-from .call_request_continue import call_req_c_rw
-from .call_response_continue import call_res_c_rw
+from .types import Types
 
 RW = {
     Types.CALL_REQ: call_req_rw,
+    Types.CALL_REQ_CONTINUE: call_req_c_rw,
     Types.CALL_RES: call_res_rw,
+    Types.CALL_RES_CONTINUE: call_res_c_rw,
+    Types.CANCEL: cancel_rw,
+    Types.CLAIM: claim_rw,
     Types.ERROR: error_rw,
     Types.INIT_REQ: init_req_rw,
     Types.INIT_RES: init_res_rw,
     Types.PING_REQ: ping_req_rw,
     Types.PING_RES: ping_res_rw,
-    Types.CALL_REQ_CONTINUE: call_req_c_rw,
-    Types.CALL_RES_CONTINUE: call_res_c_rw,
 }
 
 __all__ = [
@@ -51,6 +55,8 @@
     "CallRequestContinueMessage",
     "CallResponseMessage",
     "CallResponseContinueMessage",
+    "CancelMessage",
+    "ClaimMessage",
     "ErrorMessage",
     "ErrorCode",
     "InitRequestMessage",

tchannel/messages/cancel.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2015 Uber Technologies, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from __future__ import absolute_import
+
+from . import common
+from .. import rw
+from ..glossary import DEFAULT_TTL
+from .base import BaseMessage
+
+
+class CancelMessage(BaseMessage):
+    __slots__ = BaseMessage.__slots__ + (
+        'ttl',
+        'tracing',
+        'why',
+    )
+
+    def __init__(self, ttl=DEFAULT_TTL, tracing=None, why=None, id=0):
+        super(CancelMessage, self).__init__(id)
+        self.ttl = ttl
+        self.tracing = tracing or common.Tracing(0, 0, 0, 0)
+        self.why = why or ''
+
+
+cancel_rw = rw.instance(
+    CancelMessage,
+    ('ttl', rw.number(4)),                          # ttl:4
+    ('tracing', common.tracing_rw),                 # tracing:24
+    ('why', rw.len_prefixed_string(rw.number(2))),  # why:2
+)

tchannel/messages/claim.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2015 Uber Technologies, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from __future__ import absolute_import
+
+from . import common
+from .. import rw
+from ..glossary import DEFAULT_TTL
+from .base import BaseMessage
+
+
+class ClaimMessage(BaseMessage):
+    __slots__ = BaseMessage.__slots__ + (
+        'ttl',
+        'tracing',
+    )
+
+    def __init__(self, ttl=DEFAULT_TTL, tracing=None, id=0):
+        super(ClaimMessage, self).__init__(id)
+        self.ttl = ttl
+        self.tracing = tracing or common.Tracing(0, 0, 0, 0)
+
+
+claim_rw = rw.instance(
+    ClaimMessage,
+    ('ttl', rw.number(4)),           # ttl:4
+    ('tracing', common.tracing_rw),  # tracing:24
+)

tests/test_messages.py

@@ -152,6 +152,15 @@ def test_valid_ping_request():
         'checksum': (messages.ChecksumType.crc32, 1),
         'args': None,
     }),
+    (messages.ClaimMessage, messages.claim_rw, {
+        'ttl': 4,
+        'tracing': messages.Tracing(0, 0, 0, 1),
+    }),
+    (messages.CancelMessage, messages.cancel_rw, {
+        'ttl': 4,
+        'tracing': messages.Tracing(0, 0, 0, 1),
+        'why': 'foo',
+    }),
 ])
 def test_roundtrip_message(message_class, message_rw, attrs):
     """Verify all message types serialize and deserialize properly."""