Using topo proto3 in tests to get tablet types.

py/vtdb/keyrange_constants.py

@@ -27,17 +27,15 @@ PROTO3_KIT_TO_STRING = {
 # (Eventually we will use the proto3 version of this)
 PROTO3_TABLET_TYPE_TO_STRING = {
   0: 'unknown',
-  1: 'idle',
-  2: 'master',
-  3: 'replica',
-  4: 'rdonly',
-  5: 'spare',
-  6: 'experimental',
-  7: 'schema_upgrade',
-  8: 'backup',
-  9: 'restore',
-  10: 'worker',
-  11: 'scrap',
+  1: 'master',
+  2: 'replica',
+  3: 'rdonly',
+  4: 'spare',
+  5: 'experimental',
+  6: 'backup',
+  7: 'restore',
+  8: 'worker',
+  9: 'scrap',
 }
 
 # Converts a bson-encoded proto3 SrvKeyspace into the format

test/reparent.py

@@ -9,6 +9,8 @@ import logging
 import time
 import unittest
 
+from vtproto import topodata_pb2
+
 import environment
 import utils
 import tablet
@@ -546,8 +548,7 @@ class TestReparent(unittest.TestCase):
     health = utils.run_vtctl_json(['VtTabletStreamHealth',
                                    '-count', '1',
                                    tablet_62044.tablet_alias])
-    self.assertEqual(health['target']['tablet_type'],
-                     tablet.Tablet.tablet_type_value['MASTER'])
+    self.assertEqual(health['target']['tablet_type'], topodata_pb2.MASTER)
     # have to compare the int version, or the rounding errors can break
     self.assertTrue(
         health['tablet_externally_reparented_timestamp'] >= int(base_time))

test/resharding.py

@@ -10,6 +10,8 @@ import threading
 import time
 import unittest
 
+from vtproto import topodata_pb2
+
 from vtdb import keyrange_constants
 
 import environment
@@ -703,8 +705,7 @@ primary key (name)
     # Destination tablets would have query service disabled for other reasons than the migration,
     # so check the shard record instead of the tablets directly
     utils.check_shard_query_services(self, destination_shards,
-                                     tablet.Tablet.tablet_type_value['REPLICA'],
-                                     False)
+                                     topodata_pb2.REPLICA, False)
     utils.check_srv_keyspace('test_nj', 'test_keyspace',
                              'Partitions(master): -80 80-\n'
                              'Partitions(rdonly): -80 80-c0 c0-\n'
@@ -720,8 +721,7 @@ primary key (name)
     # reasons than the migration, so check the shard record instead of
     # the tablets directly
     utils.check_shard_query_services(self, destination_shards,
-                                     tablet.Tablet.tablet_type_value['REPLICA'],
-                                     True)
+                                     topodata_pb2.REPLICA, True)
     utils.check_srv_keyspace('test_nj', 'test_keyspace',
                              'Partitions(master): -80 80-\n'
                              'Partitions(rdonly): -80 80-c0 c0-\n'

test/tabletmanager.py

@@ -12,6 +12,8 @@ import unittest
 import urllib
 import urllib2
 
+from vtproto import topodata_pb2
+
 import environment
 import utils
 import tablet
@@ -352,7 +354,7 @@ class TestTabletManager(unittest.TestCase):
         tablet.get_healthz()
 
   def wait_for_tablet_type_change(self, tablet_alias, expected_type):
-    t = tablet.Tablet.tablet_type_value[expected_type.upper()]
+    t = topodata_pb2.TabletType.Value(expected_type.upper())
     timeout = 10
     while True:
       ti = utils.run_vtctl_json(['GetTablet', tablet_alias])
@@ -390,7 +392,7 @@ class TestTabletManager(unittest.TestCase):
 
     # make sure the master is still master
     ti = utils.run_vtctl_json(['GetTablet', tablet_62344.tablet_alias])
-    self.assertEqual(ti['type'], tablet.Tablet.tablet_type_value['MASTER'],
+    self.assertEqual(ti['type'], topodata_pb2.MASTER,
                      'unexpected master type: %s' % ti['type'])
 
     # stop replication, make sure we go unhealthy.
@@ -453,7 +455,7 @@ class TestTabletManager(unittest.TestCase):
       self.assertNotIn('tablet_externally_reparented_timestamp', data)
       self.assertEqual('test_keyspace', data['target']['keyspace'])
       self.assertEqual('0', data['target']['shard'])
-      self.assertEqual(3, data['target']['tablet_type'])
+      self.assertEqual(topodata_pb2.REPLICA, data['target']['tablet_type'])
 
     # kill the tablets
     tablet.kill_tablets([tablet_62344, tablet_62044])
@@ -462,7 +464,7 @@ class TestTabletManager(unittest.TestCase):
     # to reset its state to spare
     ti = utils.run_vtctl_json(['GetTablet', tablet_62044.tablet_alias])
     self.assertEqual(
-        ti['type'], tablet.Tablet.tablet_type_value['SPARE'],
+        ti['type'], topodata_pb2.SPARE,
         "tablet didn't go to spare while in lameduck mode: %s" % str(ti))
 
     # Also the replica should be gone from the serving graph.

test/update_stream.py

@@ -10,6 +10,8 @@ import traceback
 import threading
 import unittest
 
+from vtproto import topodata_pb2
+
 import environment
 import tablet
 import utils
@@ -167,7 +169,7 @@ class TestUpdateStream(unittest.TestCase):
     self._exec_vt_txn(['delete from vt_insert_test'])
     utils.run_vtctl(['ChangeSlaveType', replica_tablet.tablet_alias, 'spare'])
     utils.wait_for_tablet_type(
-        replica_tablet.tablet_alias, tablet.Tablet.tablet_type_value['SPARE'])
+        replica_tablet.tablet_alias, topodata_pb2.SPARE)
     logging.debug('dialing replica update stream service')
     replica_conn = self._get_replica_stream_conn()
     try:
@@ -193,9 +195,8 @@ class TestUpdateStream(unittest.TestCase):
     utils.run_vtctl(
         ['ChangeSlaveType', replica_tablet.tablet_alias, 'replica'])
     logging.debug('sleeping a bit for the replica action to complete')
-    utils.wait_for_tablet_type(
-        replica_tablet.tablet_alias,
-        tablet.Tablet.tablet_type_value['REPLICA'], 30)
+    utils.wait_for_tablet_type(replica_tablet.tablet_alias,
+                               topodata_pb2.REPLICA, 30)
     thd = threading.Thread(target=self.perform_writes, name='write_thd',
                            args=(100,))
     thd.daemon = True
@@ -230,9 +231,8 @@ class TestUpdateStream(unittest.TestCase):
         if first:
           utils.run_vtctl(
               ['ChangeSlaveType', replica_tablet.tablet_alias, 'spare'])
-          utils.wait_for_tablet_type(
-              replica_tablet.tablet_alias,
-              tablet.Tablet.tablet_type_value['SPARE'], 30)
+          utils.wait_for_tablet_type(replica_tablet.tablet_alias,
+                                     topodata_pb2.SPARE, 30)
           first = False
         else:
           if stream_event.category == update_stream.StreamEvent.POS:
@@ -362,9 +362,8 @@ class TestUpdateStream(unittest.TestCase):
     # The above tests leaves the service in disabled state, hence enabling it.
     utils.run_vtctl(
         ['ChangeSlaveType', replica_tablet.tablet_alias, 'replica'])
-    utils.wait_for_tablet_type(
-        replica_tablet.tablet_alias,
-        tablet.Tablet.tablet_type_value['REPLICA'], 30)
+    utils.wait_for_tablet_type(replica_tablet.tablet_alias,
+                               topodata_pb2.REPLICA, 30)
 
   def test_log_rotation(self):
     start_position = _get_master_current_position()

test/utils.py

@@ -17,16 +17,20 @@ import urllib2
 
 import MySQLdb
 
-import environment
+from vtproto import topodata_pb2
 
 from vtdb import keyrange_constants
+
 from vtctl import vtctl_client
-from mysql_flavor import set_mysql_flavor
+
+import environment
 from mysql_flavor import mysql_flavor
-from protocols_flavor import set_protocols_flavor
+from mysql_flavor import set_mysql_flavor
 from protocols_flavor import protocols_flavor
+from protocols_flavor import set_protocols_flavor
 from topo_flavor.server import set_topo_server_flavor
 
+
 options = None
 devnull = open('/dev/null', 'w')
 hostname = socket.getaddrinfo(
@@ -869,7 +873,9 @@ def check_srv_keyspace(cell, keyspace, expected, keyspace_id_type='uint64'):
   result = ''
   pmap = {}
   for partition in ks['partitions']:
-    tablet_type = keyrange_constants.PROTO3_TABLET_TYPE_TO_STRING[partition['served_type']]
+    tablet_type = topodata_pb2.TabletType.Name(partition['served_type']).lower()
+    if tablet_type == 'batch':
+      tablet_type = 'rdonly'
     r = 'Partitions(%s):' % tablet_type
     for shard in partition['shard_references']:
       s = ''

test/vertical_split.py

@@ -8,6 +8,8 @@ import logging
 import time
 import unittest
 
+from vtproto import topodata_pb2
+
 from vtdb import keyrange
 from vtdb import keyrange_constants
 from vtdb import vtgate_client
@@ -170,7 +172,9 @@ index by_msg (msg)
     if 'served_from' in ks and ks['served_from']:
       a = []
       for served_from in sorted(ks['served_from']):
-        tt = keyrange_constants.PROTO3_TABLET_TYPE_TO_STRING[served_from['tablet_type']]
+        tt = topodata_pb2.TabletType.Name(served_from['tablet_type']).lower()
+        if tt == 'batch':
+          tt = 'rdonly'
         a.append('ServedFrom(%s): %s\n' % (tt, served_from['keyspace']))
       for line in sorted(a):
         result += line
@@ -427,7 +431,7 @@ index by_msg (msg)
         ['GetKeyspace', 'destination_keyspace'])
     found = False
     for ksf in keyspace_json['served_froms']:
-      if ksf['tablet_type'] == 4:
+      if ksf['tablet_type'] == topodata_pb2.RDONLY:
         found = True
         self.assertEqual(ksf['cells'], ['test_nj'])
     self.assertTrue(found)
@@ -438,7 +442,7 @@ index by_msg (msg)
         ['GetKeyspace', 'destination_keyspace'])
     found = False
     for ksf in keyspace_json['served_froms']:
-      if ksf['tablet_type'] == 4:
+      if ksf['tablet_type'] == topodata_pb2.RDONLY:
         found = True
     self.assertFalse(found)
     utils.run_vtctl(['SetKeyspaceServedFrom', '-source=source_keyspace',
@@ -448,7 +452,7 @@ index by_msg (msg)
         ['GetKeyspace', 'destination_keyspace'])
     found = False
     for ksf in keyspace_json['served_froms']:
-      if ksf['tablet_type'] == 4:
+      if ksf['tablet_type'] == topodata_pb2.RDONLY:
         found = True
         self.assertNotIn('cells', ksf)
     self.assertTrue(found)
@@ -517,14 +521,13 @@ index by_msg (msg)
     shard_json = utils.run_vtctl_json(['GetShard', 'source_keyspace/0'])
     self.assertEqual(len(shard_json['tablet_controls']), 2)
     for tc in shard_json['tablet_controls']:
-      self.assertIn(tc['tablet_type'], [
-          tablet.Tablet.tablet_type_value['MASTER'],
-          tablet.Tablet.tablet_type_value['REPLICA']])
+      self.assertIn(tc['tablet_type'], [topodata_pb2.MASTER,
+                                        topodata_pb2.REPLICA])
     utils.run_vtctl(['SetShardTabletControl', '--tables=moving.*,view1',
                      'source_keyspace/0', 'rdonly'], auto_log=True)
     shard_json = utils.run_vtctl_json(['GetShard', 'source_keyspace/0'])
     for tc in shard_json['tablet_controls']:
-      if tc['tablet_type'] == 4:
+      if tc['tablet_type'] == topodata_pb2.RDONLY:
         break
     self.assertEqual(['moving.*', 'view1'], tc['blacklisted_tables'])
     utils.run_vtctl(['SetShardTabletControl', '--remove', 'source_keyspace/0',

test/vtctld_test.py

@@ -5,6 +5,8 @@ import unittest
 import urllib2
 import re
 
+from vtproto import topodata_pb2
+
 import environment
 import tablet
 import utils
@@ -209,11 +211,10 @@ class TestVtctld(unittest.TestCase):
     for tn in s0['TabletNodes']:
       tt = tn['TabletType']
       types.append(tt)
-      if tt == tablet.Tablet.tablet_type_value['MASTER']:
+      if tt == topodata_pb2.MASTER:
         self.assertEqual(len(tn['Nodes']), 1)
-    self.assertItemsEqual(sorted(types), [
-        tablet.Tablet.tablet_type_value['MASTER'],
-        tablet.Tablet.tablet_type_value['REPLICA']])
+    self.assertItemsEqual(sorted(types),
+                          sorted([topodata_pb2.MASTER, topodata_pb2.REPLICA]))
     self.assertEqual(
         self.serving_data['redirected_keyspace']['ServedFrom']['master'],
         'test_keyspace')