From 8c6e9cca323b3b49c48668916bf646d9bda54037 Mon Sep 17 00:00:00 2001
From: Martin Raspaud <martin.raspaud@smhi.se>
Date: Fri, 16 Nov 2018 15:30:41 +0100
Subject: [PATCH] Make sure messages are consistently encode as unicode

---
 posttroll/message.py           | 20 +++++++++++++++++---
 posttroll/publisher.py         |  4 ++--
 posttroll/tests/test_pubsub.py | 19 +++++++++++++++++++
 3 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/posttroll/message.py b/posttroll/message.py
index 8beba89..7201924 100644
--- a/posttroll/message.py
+++ b/posttroll/message.py
@@ -119,7 +119,14 @@ def __init__(self, subject='', atype='', data='', binary=False, rawstr=None):
         if rawstr:
             self.__dict__ = _decode(rawstr)
         else:
-            self.subject = subject
+            try:
+                self.subject = subject.decode('utf-8')
+            except AttributeError:
+                self.subject = subject
+            try:
+                self.type = atype.decode('utf-8')
+            except AttributeError:
+                self.type = atype
             self.type = atype
             self.sender = _getsender()
             self.time = datetime.utcnow()
@@ -168,9 +175,16 @@ def encode(self):
     def __repr__(self):
         return self.encode()
 
-    def __str__(self):
+    def __unicode__(self):
         return self.encode()
 
+    def __str__(self):
+        try:
+            return unicode(self).encode('utf-8')
+        except NameError:
+            return self.encode()
+
+
     def _validate(self):
         """Validate a messages attributes.
         """
@@ -299,7 +313,7 @@ def datetime_encoder(obj):
 def _encode(msg, head=False, binary=False):
     """Convert a Message to a raw string.
     """
-    rawstr = str(_MAGICK) + "{0:s} {1:s} {2:s} {3:s} {4:s}".format(
+    rawstr = str(_MAGICK) + u"{0:s} {1:s} {2:s} {3:s} {4:s}".format(
         msg.subject, msg.type, msg.sender, msg.time.isoformat(), msg.version)
 
     if not head and msg.data:
diff --git a/posttroll/publisher.py b/posttroll/publisher.py
index ddf386b..2113da4 100644
--- a/posttroll/publisher.py
+++ b/posttroll/publisher.py
@@ -79,7 +79,7 @@ class Publisher(object):
             while True:
                 counter += 1
                 message = Message("/counter", "info", str(counter))
-                pub.send_string(str(message))
+                pub.send_string(message.encode())
                 time.sleep(3)
         except KeyboardInterrupt:
             print("terminating publisher...")
@@ -247,7 +247,7 @@ class Publish(NoisyPublisher):
                         counter += 1
                         message = Message("/counter", "info", str(counter))
                         print("publishing", message)
-                        pub.send(str(message))
+                        pub.send(message.encode())
                         time.sleep(3)
             except KeyboardInterrupt:
                 print("terminating publisher...")
diff --git a/posttroll/tests/test_pubsub.py b/posttroll/tests/test_pubsub.py
index 12b0f1f..011da79 100644
--- a/posttroll/tests/test_pubsub.py
+++ b/posttroll/tests/test_pubsub.py
@@ -236,6 +236,24 @@ def test_pub_suber(self):
         self.assertTrue(tested)
         pub.stop()
 
+class TestPub(unittest.TestCase):
+
+    """Testing the publishing capabilities.
+    """
+
+    def setUp(self):
+        test_lock.acquire()
+
+    def tearDown(self):
+        test_lock.release()
+
+    def test_pub_unicode(self):
+        message = Message("/pџтяöll", "info", 'hej')
+        with Publish("a_service", 9000) as pub:
+            try:
+                pub.send(message.encode())
+            except UnicodeDecodeError:
+                self.fail("Sending raised UnicodeDecodeError unexpectedly!")
 
 class TestListenerContainer(unittest.TestCase):
 
@@ -282,5 +300,6 @@ def suite():
     mysuite.addTest(loader.loadTestsFromTestCase(TestNS))
     mysuite.addTest(loader.loadTestsFromTestCase(TestNSWithoutMulticasting))
     mysuite.addTest(loader.loadTestsFromTestCase(TestListenerContainer))
+    mysuite.addTest(loader.loadTestsFromTestCase(TestPub))
 
     return mysuite