Added unit tests for vtdb.tablet

Makefile

@@ -83,6 +83,7 @@ site_integration_test_files = \
 # - medium: 30 secs - 1 min
 # - large: over 1 min
 small_integration_test_files = \
+	tablet_test.py \
 	vertical_split.py \
 	vertical_split_vtgate.py \
 	schema.py \

py/vtdb/tablet.py

@@ -173,9 +173,14 @@ class TabletConnection(object):
     """
     response = self.client.call(method_name, request)
     reply = response.reply
+    if not reply or not isinstance(reply, dict):
+      return response
     # Handle the case of new client => old server
-    if reply.get('Err', None):
-      raise gorpc.AppError(reply['Err'].get('Message', 'Missing error message'), method_name)
+    err = reply.get('Err', None)
+    if err:
+      if not isinstance(reply, dict) or 'Message' not in err:
+        raise gorpc.AppError('Missing error message', method_name)
+      raise gorpc.AppError(reply['Err']['Message'], method_name)
     return response
 
   def _execute(self, sql, bind_variables):

test/config.json

@@ -82,6 +82,9 @@
     },
     "worker": {
       "File": "worker.py"
+    },
+    "tablet": {
+      "File": "tablet_test.py"
     }
   }
 }

test/tablet_test.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""Unit tests for vtdb.tablet"""
+
+import unittest
+
+import mock
+
+from net import gorpc
+import utils
+from vtdb import tablet
+
+
+class TestRPCCallAndExtract(unittest.TestCase):
+  """Tests rpc_call_and_extract_error is tolerant to various responses."""
+
+  tablet_conn = tablet.TabletConnection('addr', 'type', 'keyspace', 'shard', 30)
+
+  def test_reply_is_none(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      mock_client.call.return_value = gorpc.GoRpcResponse()
+      self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+  def test_reply_is_empty_string(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      response = gorpc.GoRpcResponse()
+      response.reply = ''
+      mock_client.call.return_value = response
+      self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+  def test_reply_is_string(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      response = gorpc.GoRpcResponse()
+      response.reply = 'foo'
+      mock_client.call.return_value = response
+      self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+  def test_reply_is_dict(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      response = gorpc.GoRpcResponse()
+      response.reply = {'foo': 'bar'}
+      mock_client.call.return_value = response
+      self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+  def test_reply_has_non_dict_err(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      response = gorpc.GoRpcResponse()
+      response.reply = {'Err': 'foo'}
+      mock_client.call.return_value = response
+      with self.assertRaisesRegexp(gorpc.AppError, 'Missing error message'):
+        self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+  def test_reply_has_missing_err_message(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      response = gorpc.GoRpcResponse()
+      response.reply = {'Err': {'foo': 'bar'}}
+      mock_client.call.return_value = response
+      with self.assertRaisesRegexp(gorpc.AppError, 'Missing error message'):
+        self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+  def test_reply_has_err_message(self):
+    with mock.patch.object(self.tablet_conn, 'client', autospec=True) as mock_client:
+      response = gorpc.GoRpcResponse()
+      response.reply = {'Err': {'Message': 'bar'}}
+      mock_client.call.return_value = response
+      with self.assertRaisesRegexp(gorpc.AppError, "'bar', 'method'"):
+        self.tablet_conn.rpc_call_and_extract_error('method', 'req')
+
+if __name__ == '__main__':
+  utils.main()