From 7b12f843666d50430a96a8848e7a540ed28dffef Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 29 Mar 2018 11:53:02 +0545
Subject: [PATCH] core: support CallError(str) for service.py.

---
 mitogen/core.py          | 19 +++++++-----
 test.sh                  |  1 +
 tests/call_error_test.py | 66 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 7 deletions(-)
 create mode 100644 tests/call_error_test.py

diff --git a/mitogen/core.py b/mitogen/core.py
index e9026f03..4cac056c 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -95,13 +95,18 @@ class LatchError(Error):
 
 
 class CallError(Error):
-    def __init__(self, e):
-        s = '%s.%s: %s' % (type(e).__module__, type(e).__name__, e)
-        tb = sys.exc_info()[2]
-        if tb:
-            s += '\n'
-            s += ''.join(traceback.format_tb(tb))
-        Error.__init__(self, s)
+    def __init__(self, fmt=None, *args):
+        if not isinstance(fmt, Exception):
+            Error.__init__(self, fmt, *args)
+        else:
+            e = fmt
+            fmt = '%s.%s: %s' % (type(e).__module__, type(e).__name__, e)
+            args = ()
+            tb = sys.exc_info()[2]
+            if tb:
+                fmt += '\n'
+                fmt += ''.join(traceback.format_tb(tb))
+            Error.__init__(self, fmt)
 
     def __reduce__(self):
         return (_unpickle_call_error, (self[0],))
diff --git a/test.sh b/test.sh
index 6e678674..f254fcf9 100755
--- a/test.sh
+++ b/test.sh
@@ -36,6 +36,7 @@ run_test()
 }
 
 run_test tests/ansible_helpers_test.py
+run_test tests/call_error_test.py
 run_test tests/call_function_test.py
 run_test tests/channel_test.py
 run_test tests/fakessh_test.py
diff --git a/tests/call_error_test.py b/tests/call_error_test.py
new file mode 100644
index 00000000..f96aae27
--- /dev/null
+++ b/tests/call_error_test.py
@@ -0,0 +1,66 @@
+import os
+import pickle
+
+import unittest2
+
+import mitogen.core
+
+
+class ConstructorTest(unittest2.TestCase):
+    klass = mitogen.core.CallError
+
+    def test_string_noargs(self):
+        e = self.klass('%s%s')
+        self.assertEquals(e[0], '%s%s')
+
+    def test_string_args(self):
+        e = self.klass('%s%s', 1, 1)
+        self.assertEquals(e[0], '11')
+
+    def test_from_exc(self):
+        ve = ValueError('eek')
+        e = self.klass(ve)
+        self.assertEquals(e[0], 'exceptions.ValueError: eek')
+
+    def test_from_exc_tb(self):
+        try:
+            raise ValueError('eek')
+        except ValueError, ve:
+            e = self.klass(ve)
+
+        self.assertTrue(e[0].startswith('exceptions.ValueError: eek'))
+        self.assertTrue('test_from_exc_tb' in e[0])
+
+
+class PickleTest(unittest2.TestCase):
+    klass = mitogen.core.CallError
+
+    def test_string_noargs(self):
+        e = self.klass('%s%s')
+        e2 = pickle.loads(pickle.dumps(e))
+        self.assertEquals(e2[0], '%s%s')
+
+    def test_string_args(self):
+        e = self.klass('%s%s', 1, 1)
+        e2 = pickle.loads(pickle.dumps(e))
+        self.assertEquals(e2[0], '11')
+
+    def test_from_exc(self):
+        ve = ValueError('eek')
+        e = self.klass(ve)
+        e2 = pickle.loads(pickle.dumps(e))
+        self.assertEquals(e2[0], 'exceptions.ValueError: eek')
+
+    def test_from_exc_tb(self):
+        try:
+            raise ValueError('eek')
+        except ValueError, ve:
+            e = self.klass(ve)
+
+        e2 = pickle.loads(pickle.dumps(e))
+        self.assertTrue(e2[0].startswith('exceptions.ValueError: eek'))
+        self.assertTrue('test_from_exc_tb' in e2[0])
+
+
+if __name__ == '__main__':
+    unittest2.main()