From 3da64bc65dccd1739b03d1f3c5c34e8e3ba9411b Mon Sep 17 00:00:00 2001
From: Nicolas Jouanin <nico@beerfactory.org>
Date: Wed, 17 Jun 2015 17:28:36 +0200
Subject: [PATCH] Add PUBCOMP, PUBREL messages

---
 hbmqtt/mqtt/pubcomp.py     | 38 ++++++++++++++++++++++++++++++++++++++
 hbmqtt/mqtt/pubrel.py      | 38 ++++++++++++++++++++++++++++++++++++++
 tests/mqtt/test_pubcomp.py | 24 ++++++++++++++++++++++++
 tests/mqtt/test_pubrel.py  | 24 ++++++++++++++++++++++++
 4 files changed, 124 insertions(+)
 create mode 100644 hbmqtt/mqtt/pubcomp.py
 create mode 100644 hbmqtt/mqtt/pubrel.py
 create mode 100644 tests/mqtt/test_pubcomp.py
 create mode 100644 tests/mqtt/test_pubrel.py

diff --git a/hbmqtt/mqtt/pubcomp.py b/hbmqtt/mqtt/pubcomp.py
new file mode 100644
index 0000000..b82d4f8
--- /dev/null
+++ b/hbmqtt/mqtt/pubcomp.py
@@ -0,0 +1,38 @@
+# Copyright (c) 2015 Nicolas JOUANIN
+#
+# See the file license.txt for copying permission.
+from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.errors import HBMQTTException
+from hbmqtt.codecs import *
+
+
+class PubcompVariableHeader(MQTTVariableHeader):
+    def __init__(self, packet_id):
+        super().__init__()
+        self.packet_id = packet_id
+
+    def to_bytes(self):
+        out = b''
+        out += int_to_bytes(self.packet_id, 2)
+        return out
+
+    @classmethod
+    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
+        packet_id = yield from decode_packet_id(reader)
+        return cls(packet_id)
+
+
+class PubcompPacket(MQTTPacket):
+    VARIABLE_HEADER = PubcompVariableHeader
+    PAYLOAD = None
+
+    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PubcompVariableHeader=None, payload=None):
+        if fixed is None:
+            header = MQTTFixedHeader(PacketType.PUBCOMP, 0x00)
+        else:
+            if fixed.packet_type is not PacketType.PUBCOMP:
+                raise HBMQTTException("Invalid fixed packet type %s for PubcompPacket init" % fixed.packet_type)
+            header = fixed
+        super().__init__(header)
+        self.variable_header = variable_header
+        self.payload = None
diff --git a/hbmqtt/mqtt/pubrel.py b/hbmqtt/mqtt/pubrel.py
new file mode 100644
index 0000000..da9c06c
--- /dev/null
+++ b/hbmqtt/mqtt/pubrel.py
@@ -0,0 +1,38 @@
+# Copyright (c) 2015 Nicolas JOUANIN
+#
+# See the file license.txt for copying permission.
+from hbmqtt.mqtt.packet import MQTTPacket, MQTTFixedHeader, PacketType, MQTTVariableHeader
+from hbmqtt.errors import HBMQTTException
+from hbmqtt.codecs import *
+
+
+class PubrelVariableHeader(MQTTVariableHeader):
+    def __init__(self, packet_id):
+        super().__init__()
+        self.packet_id = packet_id
+
+    def to_bytes(self):
+        out = b''
+        out += int_to_bytes(self.packet_id, 2)
+        return out
+
+    @classmethod
+    def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader):
+        packet_id = yield from decode_packet_id(reader)
+        return cls(packet_id)
+
+
+class PubrelPacket(MQTTPacket):
+    VARIABLE_HEADER = PubrelVariableHeader
+    PAYLOAD = None
+
+    def __init__(self, fixed: MQTTFixedHeader=None, variable_header: PubrelVariableHeader=None, payload=None):
+        if fixed is None:
+            header = MQTTFixedHeader(PacketType.PUBREL, 0x00)
+        else:
+            if fixed.packet_type is not PacketType.PUBREL:
+                raise HBMQTTException("Invalid fixed packet type %s for PubrelPacket init" % fixed.packet_type)
+            header = fixed
+        super().__init__(header)
+        self.variable_header = variable_header
+        self.payload = None
diff --git a/tests/mqtt/test_pubcomp.py b/tests/mqtt/test_pubcomp.py
new file mode 100644
index 0000000..94e76a3
--- /dev/null
+++ b/tests/mqtt/test_pubcomp.py
@@ -0,0 +1,24 @@
+# Copyright (c) 2015 Nicolas JOUANIN
+#
+# See the file license.txt for copying permission.
+import unittest
+
+from hbmqtt.mqtt.pubcomp import PubcompPacket, PubcompVariableHeader
+from hbmqtt.codecs import *
+
+class PubcompPacketTest(unittest.TestCase):
+    def setUp(self):
+        self.loop = asyncio.new_event_loop()
+
+    def test_from_stream(self):
+        data = b'\x70\x02\x00\x0a'
+        stream = asyncio.StreamReader(loop=self.loop)
+        stream.feed_data(data)
+        message = self.loop.run_until_complete(PubcompPacket.from_stream(stream))
+        self.assertEqual(message.variable_header.packet_id, 10)
+
+    def test_to_bytes(self):
+        variable_header = PubcompVariableHeader(10)
+        publish = PubcompPacket(variable_header=variable_header)
+        out = publish.to_bytes()
+        self.assertEqual(out, b'\x70\x02\x00\x0a')
\ No newline at end of file
diff --git a/tests/mqtt/test_pubrel.py b/tests/mqtt/test_pubrel.py
new file mode 100644
index 0000000..51225d1
--- /dev/null
+++ b/tests/mqtt/test_pubrel.py
@@ -0,0 +1,24 @@
+# Copyright (c) 2015 Nicolas JOUANIN
+#
+# See the file license.txt for copying permission.
+import unittest
+
+from hbmqtt.mqtt.pubrel import PubrelPacket, PubrelVariableHeader
+from hbmqtt.codecs import *
+
+class PubrelPacketTest(unittest.TestCase):
+    def setUp(self):
+        self.loop = asyncio.new_event_loop()
+
+    def test_from_stream(self):
+        data = b'\x60\x02\x00\x0a'
+        stream = asyncio.StreamReader(loop=self.loop)
+        stream.feed_data(data)
+        message = self.loop.run_until_complete(PubrelPacket.from_stream(stream))
+        self.assertEqual(message.variable_header.packet_id, 10)
+
+    def test_to_bytes(self):
+        variable_header = PubrelVariableHeader(10)
+        publish = PubrelPacket(variable_header=variable_header)
+        out = publish.to_bytes()
+        self.assertEqual(out, b'\x60\x02\x00\x0a')
\ No newline at end of file