Преглед на файлове

HUE-8737 [thrift] Adding test for retry paths of thrift_util.SuperClient

Romain преди 6 години
родител
ревизия
ebfa758ad3
променени са 2 файла, в които са добавени 41 реда и са изтрити 9 реда
  1. 1 2
      desktop/core/src/desktop/lib/thrift_util.py
  2. 40 7
      desktop/core/src/desktop/lib/thrift_util_test.py

+ 1 - 2
desktop/core/src/desktop/lib/thrift_util.py

@@ -36,8 +36,7 @@ import sys
 
 from thrift.Thrift import TType, TApplicationException
 from thrift.transport.TSocket import TSocket
-from thrift.transport.TTransport import TBufferedTransport, TFramedTransport, TMemoryBuffer,\
-                                        TTransportException
+from thrift.transport.TTransport import TBufferedTransport, TFramedTransport, TMemoryBuffer, TTransportException
 from thrift.protocol.TBinaryProtocol import TBinaryProtocol
 from thrift.protocol.TMultiplexedProtocol import TMultiplexedProtocol
 

+ 40 - 7
desktop/core/src/desktop/lib/thrift_util_test.py

@@ -25,24 +25,28 @@ import threading
 import time
 import unittest
 
-LOG = logging.getLogger(__name__)
+if sys.version_info[0] > 2:
+  from unittest.mock import patch, Mock
+else:
+  from mock import patch, Mock
 
 gen_py_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "gen-py"))
 if not gen_py_path in sys.path:
   sys.path.insert(1, gen_py_path)
 
-from djangothrift_test_gen.ttypes import TestStruct, TestNesting, TestEnum, TestManyTypes
 from djangothrift_test_gen import TestService
+from djangothrift_test_gen.ttypes import TestStruct, TestNesting, TestEnum, TestManyTypes
+from nose.tools import assert_equal, assert_true
+from thrift.protocol.TBinaryProtocol import TBinaryProtocolFactory
+from thrift.server import TServer
+from thrift.transport import TSocket
+from thrift.transport.TTransport import TBufferedTransportFactory, TTransportException
 
 from desktop.lib import python_util, thrift_util
 from desktop.lib.thrift_util import jsonable2thrift, thrift2json, _unpack_guid_secret_in_handle
 
-from thrift.protocol.TBinaryProtocol import TBinaryProtocolFactory
-from thrift.server import TServer
-from thrift.transport import TSocket
-from thrift.transport.TTransport import TBufferedTransportFactory
 
-from nose.tools import assert_equal
+LOG = logging.getLogger(__name__)
 
 
 class SimpleThriftServer(object):
@@ -303,5 +307,34 @@ class TestJsonable2Thrift(unittest.TestCase):
     self.assertBackAndForth(TestManyTypes(a_string_list=["alpha", "beta"]))
     self.assertBackAndForth(TestManyTypes(a_string_list=[u"alpha", u"beta"]))
 
+
+class TestSuperClient(unittest.TestCase):
+
+  def test_wrapper_no_retry(self):
+    wrapped_client, transport = Mock(), Mock()
+    wrapped_client.my_call = Mock(
+      side_effect=TTransportException(message='read operation timed out')
+    )
+
+    client = thrift_util.SuperClient(wrapped_client, transport)
+
+    with self.assertRaises(TTransportException):
+      client.my_call()
+      # Could check output for "Not retrying thrift call my_call due to socket timeout"
+
+
+  def test_wrapper_with_retry(self):
+    wrapped_client, transport = Mock(), Mock()
+    wrapped_client.my_call = Mock(
+      side_effect=TTransportException(message='some error')
+    )
+
+    client = thrift_util.SuperClient(wrapped_client, transport)
+
+    with self.assertRaises(TTransportException):
+      client.my_call()
+      # Could check output for several "Thrift exception; retrying: some error"
+
+
 if __name__ == '__main__':
   unittest.main()