Make topo server flavor pluggable in integration tests.

test/binlog.py

@@ -28,7 +28,7 @@ dst_replica = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         src_master.init_mysql(),
@@ -118,7 +118,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/clone.py

@@ -23,7 +23,7 @@ tablet_31981 = tablet.Tablet(31981, use_mysqlctld=use_mysqlctld)
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     # start mysql instance external to the test
     global setup_procs
@@ -55,7 +55,7 @@ def tearDownModule():
         ]
     utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 
@@ -71,7 +71,7 @@ def tearDownModule():
 class TestClone(unittest.TestCase):
   def tearDown(self):
     tablet.Tablet.check_vttablet_count()
-    environment.topo_server_wipe()
+    environment.topo_server().wipe()
     for t in [tablet_62344, tablet_31981]:
       t.reset_replication()
       t.clean_dbs()

test/environment.py

@@ -6,6 +6,12 @@ import os
 import socket
 import subprocess
 
+# Import the topo implementations that you want registered as options for the
+# --topo-server-flavor flag.
+import topo_flavor.zookeeper
+
+from topo_flavor.server import topo_server
+
 # vttop is the toplevel of the vitess source tree
 vttop = os.environ['VTTOP']
 
@@ -114,58 +120,3 @@ def binary_argstr(name):
 # binary management for the MySQL distribution.
 def mysql_binary_path(name):
   return os.path.join(vt_mysql_root, 'bin', name)
-
-# topology server management: we use zookeeper in all the tests
-topo_server_implementation = 'zookeeper'
-hostname = socket.gethostname()
-zk_port_base = reserve_ports(3)
-zkocc_port_base = reserve_ports(3)
-def topo_server_setup(add_bad_host=False):
-  global zk_port_base
-  global zkocc_port_base
-  zk_ports = ":".join([str(zk_port_base), str(zk_port_base+1), str(zk_port_base+2)])
-  run(binary_args('zkctl') + [
-       '-log_dir', vtlogroot,
-       '-zk.cfg', '1@%s:%s' % (hostname, zk_ports),
-       'init'])
-  config = tmproot+'/test-zk-client-conf.json'
-  with open(config, 'w') as f:
-    ca_server = 'localhost:%u' % (zk_port_base+2)
-    if add_bad_host:
-      ca_server += ',does.not.exists:1234'
-    zk_cell_mapping = {'test_nj': 'localhost:%u'%(zk_port_base+2),
-                       'test_ny': 'localhost:%u'%(zk_port_base+2),
-                       'test_ca': ca_server,
-                       'global': 'localhost:%u'%(zk_port_base+2),
-                       'test_nj:_zkocc': 'localhost:%u,localhost:%u,localhost:%u'%(zkocc_port_base,zkocc_port_base+1,zkocc_port_base+2),
-                       'test_ny:_zkocc': 'localhost:%u'%(zkocc_port_base),
-                       'test_ca:_zkocc': 'localhost:%u'%(zkocc_port_base),
-                       'global:_zkocc': 'localhost:%u'%(zkocc_port_base),}
-    json.dump(zk_cell_mapping, f)
-  os.environ['ZK_CLIENT_CONFIG'] = config
-  run(binary_args('zk') + ['touch', '-p', '/zk/test_nj/vt'])
-  run(binary_args('zk') + ['touch', '-p', '/zk/test_ny/vt'])
-  run(binary_args('zk') + ['touch', '-p', '/zk/test_ca/vt'])
-
-def topo_server_teardown():
-  global zk_port_base
-  import utils
-  zk_ports = ":".join([str(zk_port_base), str(zk_port_base+1), str(zk_port_base+2)])
-  run(binary_args('zkctl') + [
-       '-log_dir', vtlogroot,
-       '-zk.cfg', '1@%s:%s' % (hostname, zk_ports),
-       'shutdown' if utils.options.keep_logs else 'teardown'],
-      raise_on_error=False)
-
-def topo_server_wipe():
-  # Work around safety check on recursive delete.
-  run(binary_args('zk') + ['rm', '-rf', '/zk/test_nj/vt/*'])
-  run(binary_args('zk') + ['rm', '-rf', '/zk/test_ny/vt/*'])
-  run(binary_args('zk') + ['rm', '-rf', '/zk/global/vt/*'])
-
-  run(binary_args('zk') + ['rm', '-f', '/zk/test_nj/vt'])
-  run(binary_args('zk') + ['rm', '-f', '/zk/test_ny/vt'])
-  run(binary_args('zk') + ['rm', '-f', '/zk/global/vt'])
-
-def topo_server_flags():
-  return ['-topo_implementation', 'zookeeper']

test/initial_sharding.py

@@ -49,7 +49,7 @@ shard_1_rdonly = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         shard_master.init_mysql(),
@@ -85,7 +85,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/java_vtgate_test_helper.py

@@ -51,7 +51,7 @@ class TestEnv(object):
 
   def set_up(self):
     try:
-      environment.topo_server_setup()
+      environment.topo_server().setup()
       utils.wait_procs([t.init_mysql() for t in self.tablets])
       utils.run_vtctl(['CreateKeyspace', self.keyspace])
       utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', self.keyspace, 'keyspace_id', 'uint64'])
@@ -79,7 +79,7 @@ class TestEnv(object):
     tablet.kill_tablets(self.tablets)
     teardown_procs = [t.teardown_mysql() for t in self.tablets]
     utils.wait_procs(teardown_procs, raise_on_error=False)
-    environment.topo_server_teardown()
+    environment.topo_server().teardown()
     utils.kill_sub_processes()
     utils.remove_tmp_files()
     for t in self.tablets:

test/keyspace_test.py

@@ -64,7 +64,7 @@ primary key (id)
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         shard_0_master.init_mysql(),
@@ -105,7 +105,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/mysqlctl.py

@@ -18,7 +18,7 @@ replica_tablet = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
     utils.Vtctld().start()
 
     setup_procs = [
@@ -52,7 +52,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/primecache.py

@@ -31,7 +31,7 @@ replica = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         master.init_mysql(),
@@ -51,7 +51,7 @@ def tearDownModule():
       replica.teardown_mysql(),      ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/queryservice_test.py

@@ -18,6 +18,7 @@ from queryservice_tests import test_env
 
 from mysql_flavor import set_mysql_flavor
 from protocols_flavor import set_protocols_flavor
+from topo_flavor.server import set_topo_server_flavor
 
 
 if __name__ == "__main__":
@@ -33,12 +34,14 @@ if __name__ == "__main__":
                     help="Don't delete log files on teardown.")
   parser.add_option("--mysql-flavor", action="store", type="string")
   parser.add_option("--protocols-flavor", action="store", type="string")
+  parser.add_option("--topo-server-flavor", action="store", type="string")
   (options, args) = parser.parse_args()
 
   utils.options = options
   logging.getLogger().setLevel(logging.ERROR)
   set_mysql_flavor(options.mysql_flavor)
   set_protocols_flavor(options.protocols_flavor)
+  set_topo_server_flavor(options.topo_server_flavor)
 
   suite = unittest.TestSuite()
   if args:

test/queryservice_tests/test_env.py

@@ -186,7 +186,7 @@ class TestEnv(object):
     table_acl_config = os.path.join(environment.vttop, 'test', 'test_data', 'table_acl_config.json')
 
     if self.env == 'vttablet':
-      environment.topo_server_setup()
+      environment.topo_server().setup()
       utils.run_vtctl('CreateKeyspace -force test_keyspace')
       self.tablet.init_tablet('master', 'test_keyspace', '0')
       self.tablet.start_vttablet(
@@ -230,7 +230,7 @@ class TestEnv(object):
     if getattr(self, "txlogger", None):
       self.txlogger.terminate()
     if self.env == 'vttablet':
-      environment.topo_server_teardown()
+      environment.topo_server().teardown()
     utils.kill_sub_processes()
     utils.remove_tmp_files()
     self.tablet.remove_tree()

test/reparent.py

@@ -26,7 +26,7 @@ tablet_31981 = tablet.Tablet(31981)
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     # start mysql instance external to the test
     setup_procs = [
@@ -54,7 +54,7 @@ def tearDownModule():
   ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 
@@ -68,7 +68,7 @@ class TestReparent(unittest.TestCase):
 
   def tearDown(self):
     tablet.Tablet.check_vttablet_count()
-    environment.topo_server_wipe()
+    environment.topo_server().wipe()
     for t in [tablet_62344, tablet_62044, tablet_41983, tablet_31981]:
       t.reset_replication()
       t.clean_dbs()
@@ -284,7 +284,7 @@ class TestReparent(unittest.TestCase):
 
     # Start up a master mysql and vttablet
     tablet_62344.init_tablet('master', 'test_keyspace', shard_id, start=True)
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       shard = utils.run_vtctl_json(['GetShard', 'test_keyspace/' + shard_id])
       self.assertEqual(shard['Cells'], ['test_nj'],
                        'wrong list of cell in Shard: %s' % str(shard['Cells']))
@@ -298,7 +298,7 @@ class TestReparent(unittest.TestCase):
                              wait_for_start=False)
     for t in [tablet_62044, tablet_41983, tablet_31981]:
       t.wait_for_vttablet_state('SERVING')
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       shard = utils.run_vtctl_json(['GetShard', 'test_keyspace/' + shard_id])
       self.assertEqual(shard['Cells'], ['test_nj', 'test_ny'],
                        'wrong list of cell in Shard: %s' % str(shard['Cells']))
@@ -498,7 +498,7 @@ class TestReparent(unittest.TestCase):
     if brutal:
       tablet_62344.scrap(force=True)
       # we have some automated tools that do this too, so it's good to simulate
-      if environment.topo_server_implementation == 'zookeeper':
+      if environment.topo_server().flavor() == 'zookeeper':
         utils.run(environment.binary_args('zk') + ['rm', '-rf',
                                                    tablet_62344.zk_tablet_path])
 
@@ -520,7 +520,7 @@ class TestReparent(unittest.TestCase):
                          tablet_41983])
 
   def _test_reparent_from_outside_check(self, brutal):
-    if environment.topo_server_implementation != 'zookeeper':
+    if environment.topo_server().flavor() != 'zookeeper':
       return
     # make sure the shard replication graph is fine
     shard_replication = utils.run_vtctl_json(['GetShardReplication', 'test_nj',

test/resharding.py

@@ -49,7 +49,7 @@ shard_3_rdonly = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         shard_0_master.init_mysql(),
@@ -96,7 +96,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/rowcache_invalidator.py

@@ -38,7 +38,7 @@ def setUpModule():
   global vtgate_port
 
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     # start mysql instance external to the test
     setup_procs = [master_tablet.init_mysql(),
@@ -83,7 +83,7 @@ def tearDownModule():
                     replica_tablet.teardown_mysql()]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.vtgate_kill(vtgate_server)
   utils.kill_sub_processes()
   utils.remove_tmp_files()

test/schema.py

@@ -18,7 +18,7 @@ shard_1_replica1 = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         shard_0_master.init_mysql(),
@@ -49,7 +49,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/sharded.py

@@ -20,7 +20,7 @@ shard_1_replica = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         shard_0_master.init_mysql(),
@@ -45,7 +45,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 
@@ -118,7 +118,7 @@ class TestSharded(unittest.TestCase):
                      '-sql=' + create_vt_select_test.replace("\n", ""),
                      shard_0_replica.tablet_alias])
 
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       # start zkocc, we'll use it later, indirectly with the vtdb-zkocc driver
       zkocc_server = utils.zkocc_start()
 
@@ -175,7 +175,7 @@ class TestSharded(unittest.TestCase):
                       "2\ttest 2",
                       "10\ttest 10"],
                      driver="vtdb-streaming")
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       self._check_rows(["Index\tid\tmsg",
                         "1\ttest 1",
                         "2\ttest 2",
@@ -199,7 +199,7 @@ class TestSharded(unittest.TestCase):
 
     # make sure the schema checking works
     self._check_rows_schema_diff("vtdb")
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       self._check_rows_schema_diff("vtdb-zk")
       self._check_rows_schema_diff("vtdb-zkocc")
 
@@ -225,7 +225,7 @@ class TestSharded(unittest.TestCase):
     utils.run_vtctl(['ValidatePermissionsKeyspace', 'test_keyspace'],
                     auto_log=True)
 
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       # and create zkns on this complex keyspace, make sure a few files are created
       utils.run_vtctl(['ExportZknsForKeyspace', 'test_keyspace'])
       out, err = utils.run(environment.binary_argstr('zk')+' ls -R /zk/test_nj/zk?s/vt/test_keysp*', trap_output=True)
@@ -280,7 +280,7 @@ class TestSharded(unittest.TestCase):
         self.fail('unexpected exception: ' + str(e))
 
     utils.vtgate_kill(vtgate_server)
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       utils.kill_sub_process(zkocc_server)
     tablet.kill_tablets([shard_0_master, shard_0_replica, shard_1_master,
                          shard_1_replica])

test/tablet.py

@@ -422,7 +422,7 @@ class Tablet(object):
     """
     args = []
     args.extend(['-tablet-path', self.tablet_alias])
-    args.extend(environment.topo_server_flags())
+    args.extend(environment.topo_server().flags())
     args.extend(protocols_flavor().binlog_player_protocol_flags())
     args.extend(protocols_flavor().tablet_manager_protocol_flags())
     args.extend(['-pid_file', os.path.join(self.tablet_dir, 'vttablet.pid')])
@@ -515,7 +515,7 @@ class Tablet(object):
 
   def wait_for_vttablet_state(self, expected, timeout=60.0, port=None):
     # wait for zookeeper PID just to be sure we have it
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       if not self.checked_zk_pid:
         utils.run(environment.binary_args('zk') + ['wait', '-e', self.zk_pid],
                   stdout=utils.devnull)

test/tabletmanager.py

@@ -26,12 +26,12 @@ tablet_62044 = tablet.Tablet(62044)
 
 def setUpModule():
   try:
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       # this is a one-off test to make sure our zookeeper implementation
       # behaves with a server that is not DNS-resolveable
-      environment.topo_server_setup(add_bad_host=True)
+      environment.topo_server().setup(add_bad_host=True)
     else:
-      environment.topo_server_setup()
+      environment.topo_server().setup()
 
     # start mysql instance external to the test
     setup_procs = [
@@ -54,7 +54,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 
@@ -64,7 +64,7 @@ def tearDownModule():
 class TestTabletManager(unittest.TestCase):
   def tearDown(self):
     tablet.Tablet.check_vttablet_count()
-    environment.topo_server_wipe()
+    environment.topo_server().wipe()
     for t in [tablet_62344, tablet_62044]:
       t.reset_replication()
       t.clean_dbs()
@@ -183,7 +183,7 @@ class TestTabletManager(unittest.TestCase):
     # schedule long action in the background, sleep a little bit to make sure
     # it started to run
     args = (environment.binary_args('vtctl') +
-            environment.topo_server_flags() +
+            environment.topo_server().flags() +
             protocols_flavor().tablet_manager_protocol_flags() +
             protocols_flavor().tabletconn_protocol_flags() +
             ['-log_dir', environment.vtlogroot,
@@ -201,7 +201,7 @@ class TestTabletManager(unittest.TestCase):
     # wait for the background vtctl
     bg.wait()
 
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       # extra small test: we ran for a while, get the states we were in,
       # make sure they're accounted for properly
       # first the query engine States
@@ -312,7 +312,7 @@ class TestTabletManager(unittest.TestCase):
     tablet_62344.kill_vttablet()
 
   def test_restart(self):
-    if environment.topo_server_implementation != 'zookeeper':
+    if environment.topo_server().flavor() != 'zookeeper':
       logging.info("Skipping this test in non-github tree")
       return
 

test/topo_flavor/__init__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python
+
+# Copyright 2014, Google Inc. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can
+# be found in the LICENSE file.

test/topo_flavor/server.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+# Copyright 2014, Google Inc. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can
+# be found in the LICENSE file.
+
+import logging
+
+
+class TopoServer(object):
+  """Base class that defines the required interface."""
+
+  def setup(self):
+    """Initialize the topo server."""
+    raise NotImplementedError()
+
+  def teardown(self):
+    """Teardown the topo server."""
+    raise NotImplementedError()
+
+  def flags(self):
+    """Return a list of args that tell a Vitess process to use this topo server."""
+    raise NotImplementedError()
+
+  def wipe(self):
+    """Wipe the Vitess paths in the topo server."""
+    raise NotImplementedError()
+
+  def flavor(self):
+    """Return the name of this topo server flavor."""
+    return self.flavor_name
+
+flavor_map = {}
+
+__server = None
+
+
+def topo_server():
+  return __server
+
+
+def set_topo_server_flavor(flavor):
+  global __server
+
+  if flavor in flavor_map:
+    __server = flavor_map[flavor]
+    logging.debug("Using topo server flavor '%s'", flavor)
+  elif not flavor:
+    if len(flavor_map) == 1:
+      (flavor, __server) = flavor_map.iteritems().next()
+      logging.debug("Using default topo server flavor '%s'", flavor)
+    else:
+      logging.error(
+          "No --topo-server-flavor specified. Registered flavors: [%s]",
+          ",".join(flavor_map.keys()))
+      return
+  else:
+    logging.error(
+        "Unknown topo server flavor '%s'. Registered flavors: [%s]", flavor,
+        ",".join(flavor_map.keys()))
+    return
+
+  __server.flavor_name = flavor

test/topo_flavor/zookeeper.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+
+# Copyright 2014, Google Inc. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can
+# be found in the LICENSE file.
+
+import logging
+import os
+import socket
+import json
+
+import server
+
+
+class ZkTopoServer(server.TopoServer):
+  """Implementation of TopoServer for ZooKeeper"""
+
+  def setup(self, add_bad_host=False):
+    from environment import reserve_ports, run, binary_args, vtlogroot, tmproot
+
+    self.zk_port_base = reserve_ports(3)
+    self.zkocc_port_base = reserve_ports(3)
+
+    self.hostname = socket.gethostname()
+    self.zk_ports = ':'.join(str(self.zk_port_base + i) for i in range(3))
+    self.zk_client_port = self.zk_port_base + 2
+
+    run(binary_args('zkctl') + [
+        '-log_dir', vtlogroot,
+        '-zk.cfg', '1@%s:%s' % (self.hostname, self.zk_ports),
+        'init'])
+    config = tmproot + '/test-zk-client-conf.json'
+    with open(config, 'w') as f:
+      ca_server = 'localhost:%u' % (self.zk_client_port)
+      if add_bad_host:
+        ca_server += ',does.not.exists:1234'
+      zk_cell_mapping = {
+          'test_nj': 'localhost:%u' % (self.zk_client_port),
+          'test_ny': 'localhost:%u' % (self.zk_client_port),
+          'test_ca': ca_server,
+          'global': 'localhost:%u' % (self.zk_client_port),
+          'test_nj:_zkocc':
+              'localhost:%u,localhost:%u,localhost:%u' % tuple(
+                  self.zkocc_port_base + i
+                  for i in range(
+                      3)),
+          'test_ny:_zkocc': 'localhost:%u' % (self.zkocc_port_base),
+          'test_ca:_zkocc': 'localhost:%u' % (self.zkocc_port_base),
+          'global:_zkocc': 'localhost:%u' % (self.zkocc_port_base),
+      }
+      json.dump(zk_cell_mapping, f)
+    os.environ['ZK_CLIENT_CONFIG'] = config
+    run(binary_args('zk') + ['touch', '-p', '/zk/test_nj/vt'])
+    run(binary_args('zk') + ['touch', '-p', '/zk/test_ny/vt'])
+    run(binary_args('zk') + ['touch', '-p', '/zk/test_ca/vt'])
+
+  def teardown(self):
+    from environment import run, binary_args, vtlogroot
+    import utils
+
+    run(binary_args('zkctl') + [
+        '-log_dir', vtlogroot,
+        '-zk.cfg', '1@%s:%s' % (self.hostname, self.zk_ports),
+        'shutdown' if utils.options.keep_logs else 'teardown'],
+        raise_on_error=False)
+
+  def flags(self):
+    return ['-topo_implementation', 'zookeeper']
+
+  def wipe(self):
+    from environment import run, binary_args
+
+    # Work around safety check on recursive delete.
+    run(binary_args('zk') + ['rm', '-rf', '/zk/test_nj/vt/*'])
+    run(binary_args('zk') + ['rm', '-rf', '/zk/test_ny/vt/*'])
+    run(binary_args('zk') + ['rm', '-rf', '/zk/global/vt/*'])
+
+    run(binary_args('zk') + ['rm', '-f', '/zk/test_nj/vt'])
+    run(binary_args('zk') + ['rm', '-f', '/zk/test_ny/vt'])
+    run(binary_args('zk') + ['rm', '-f', '/zk/global/vt'])
+
+
+server.flavor_map['zookeeper'] = ZkTopoServer()

test/update_stream.py

@@ -69,7 +69,7 @@ def setUpModule():
   global master_start_position
 
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     # start mysql instance external to the test
     setup_procs = [master_tablet.init_mysql(),
@@ -131,7 +131,7 @@ def tearDownModule():
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
   utils.vtgate_kill(vtgate_server)
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
   master_tablet.remove_tree()

test/utils.py

@@ -21,6 +21,7 @@ import environment
 from vtctl import vtctl_client
 from mysql_flavor import set_mysql_flavor
 from protocols_flavor import set_protocols_flavor, protocols_flavor
+from topo_flavor.server import set_topo_server_flavor
 
 options = None
 devnull = open('/dev/null', 'w')
@@ -72,6 +73,7 @@ def main(mod=None):
   parser.add_option("-v", "--verbose", action="store_const", const=2, dest="verbose", default=1)
   parser.add_option("--mysql-flavor", action="store", type="string")
   parser.add_option("--protocols-flavor", action="store", type="string")
+  parser.add_option("--topo-server-flavor", action="store", type="string")
 
   (options, args) = parser.parse_args()
 
@@ -86,6 +88,7 @@ def main(mod=None):
 
   set_mysql_flavor(options.mysql_flavor)
   set_protocols_flavor(options.protocols_flavor)
+  set_topo_server_flavor(options.topo_server_flavor)
 
   try:
     suite = unittest.TestSuite()
@@ -383,7 +386,7 @@ def vtgate_start(vtport=None, cell='test_nj', retry_delay=1, retry_count=1,
   if topo_impl:
     args.extend(['-topo_implementation', topo_impl])
   else:
-    args.extend(environment.topo_server_flags())
+    args.extend(environment.topo_server().flags())
   if tablet_bson_encrypted:
     args.append('-tablet-bson-encrypted')
   if auth:
@@ -446,7 +449,7 @@ def run_vtctl(clargs, log_level='', auto_log=False, expect_fail=False,
 def run_vtctl_vtctl(clargs, log_level='', auto_log=False, expect_fail=False,
                     **kwargs):
   args = environment.binary_args('vtctl') + ['-log_dir', environment.vtlogroot]
-  args.extend(environment.topo_server_flags())
+  args.extend(environment.topo_server().flags())
   args.extend(protocols_flavor().tablet_manager_protocol_flags())
   args.extend(protocols_flavor().tabletconn_protocol_flags())
 
@@ -481,7 +484,7 @@ def run_vtworker(clargs, log_level='', auto_log=False, expect_fail=False, **kwar
   args = environment.binary_args('vtworker') + [
           '-log_dir', environment.vtlogroot,
           '-port', str(environment.reserve_ports(1))]
-  args.extend(environment.topo_server_flags())
+  args.extend(environment.topo_server().flags())
   args.extend(protocols_flavor().tablet_manager_protocol_flags())
 
   if auto_log:
@@ -517,7 +520,7 @@ def vtclient2(uid, path, query, bindvars=None, user=None, password=None, driver=
   server = "localhost:%u/%s" % (uid, path)
 
   cmdline = environment.binary_args('vtclient2') + ['-server', server]
-  cmdline += environment.topo_server_flags()
+  cmdline += environment.topo_server().flags()
   cmdline += protocols_flavor().tabletconn_protocol_flags()
   if user is not None:
     cmdline.extend(['-tablet-bson-username', user,
@@ -645,7 +648,7 @@ class Vtctld(object):
             '-log_dir', environment.vtlogroot,
             '-port', str(self.port),
             ] + \
-            environment.topo_server_flags() + \
+            environment.topo_server().flags() + \
             protocols_flavor().tablet_manager_protocol_flags()
     stderr_fd = open(os.path.join(environment.tmproot, "vtctld.stderr"), "w")
     self.proc = run_bg(args, stderr=stderr_fd)

test/vertical_split.py

@@ -36,7 +36,7 @@ destination_rdonly = tablet.Tablet()
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [
         source_master.init_mysql(),
@@ -68,7 +68,7 @@ def tearDownModule():
       ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 

test/vtctld_test.py

@@ -32,7 +32,7 @@ vtgate_port = None
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     setup_procs = [t.init_mysql() for t in tablets]
     utils.Vtctld().start()
@@ -52,7 +52,7 @@ def tearDownModule():
   teardown_procs = [t.teardown_mysql() for t in tablets]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 
@@ -171,9 +171,9 @@ class TestVtctld(unittest.TestCase):
     self.assertEqual(self.data["Partial"], True)
 
   def test_explorer_redirects(self):
-    if environment.topo_server_implementation != 'zookeeper':
+    if environment.topo_server().flavor() != 'zookeeper':
       logging.info('Skipping zookeeper tests in topology %s',
-                   environment.topo_server_implementation)
+                   environment.topo_server().flavor())
       return
 
     base = 'http://localhost:%u' % utils.vtctld.port

test/vtdb_test.py

@@ -70,7 +70,7 @@ primary key(eid, id)
 def setUpModule():
   logging.debug("in setUpModule")
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     # start mysql instance external to the test
     setup_procs = [shard_0_master.init_mysql(),
@@ -100,7 +100,7 @@ def tearDownModule():
                    ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
 
   utils.kill_sub_processes()
   utils.remove_tmp_files()

test/vtgatev2_test.py

@@ -72,7 +72,7 @@ pack_kid = struct.Struct('!Q').pack
 def setUpModule():
   logging.debug("in setUpModule")
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
     # start mysql instance external to the test
     setup_procs = [shard_0_master.init_mysql(),
@@ -102,7 +102,7 @@ def tearDownModule():
                    ]
   utils.wait_procs(teardown_procs, raise_on_error=False)
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
 
   utils.kill_sub_processes()
   utils.remove_tmp_files()

test/zkocc_test.py

@@ -23,7 +23,7 @@ MIN_QPS = 5000
 
 def setUpModule():
   try:
-    environment.topo_server_setup()
+    environment.topo_server().setup()
 
   except:
     tearDownModule()
@@ -34,16 +34,16 @@ def tearDownModule():
   if utils.options.skip_teardown:
     return
 
-  environment.topo_server_teardown()
+  environment.topo_server().teardown()
   utils.kill_sub_processes()
   utils.remove_tmp_files()
 
 
 class TopoOccTest(unittest.TestCase):
   def setUp(self):
-    environment.topo_server_wipe()
+    environment.topo_server().wipe()
     self.vtgate_zk, self.vtgate_zk_port = utils.vtgate_start()
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       self.zkocc_server = utils.zkocc_start()
       self.vtgate_zkocc, self.vtgate_zkocc_port = utils.vtgate_start(topo_impl="zkocc")
       self.topo = zkocc.ZkOccConnection("localhost:%u" % environment.zkocc_port_base, 'test_nj', 30)
@@ -51,7 +51,7 @@ class TopoOccTest(unittest.TestCase):
 
   def tearDown(self):
     utils.vtgate_kill(self.vtgate_zk)
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       self.topo.close()
       utils.zkocc_kill(self.zkocc_server)
       utils.vtgate_kill(self.vtgate_zkocc)
@@ -75,7 +75,7 @@ class TopoOccTest(unittest.TestCase):
     self.assertEqual(err, "KeyspaceNames[0] = test_keyspace1\n" +
                           "KeyspaceNames[1] = test_keyspace2\n")
 
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       self.assertItemsEqual(self.topo.get_srv_keyspace_names('local'), ["test_keyspace1", "test_keyspace2"])
 
       # zkocc API test
@@ -108,7 +108,7 @@ class TopoOccTest(unittest.TestCase):
                      "TabletTypes[0] = master\n",
                      "Got wrong content: %s" % err)
 
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       reply = self.topo.get_srv_keyspace("test_nj", "test_keyspace")
       self.assertEqual(reply['TabletTypes'], ['master'])
 
@@ -158,7 +158,7 @@ class TopoOccTest(unittest.TestCase):
     out, err = utils.run(environment.binary_argstr('zkclient2')+' -server localhost:%u -mode getEndPoints test_nj test_keyspace 0 master' % self.vtgate_zk_port, trap_output=True)
     self.assertEqual(err, "Entries[0] = 1 localhost\n")
 
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       self.assertEqual(len(self.topo.get_end_points("test_nj", "test_keyspace", "0", "master")['Entries']), 1)
 
       # zkocc API test
@@ -179,8 +179,8 @@ def _format_time(timeFromBson):
 class TestZkocc(unittest.TestCase):
   longMessage = True
   def setUp(self):
-    environment.topo_server_wipe()
+    environment.topo_server().wipe()
-    if environment.topo_server_implementation == 'zookeeper':
+    if environment.topo_server().flavor() == 'zookeeper':
       utils.run(environment.binary_argstr('zk')+' touch -p /zk/test_nj/vt/zkocc1')
       utils.run(environment.binary_argstr('zk')+' touch -p /zk/test_nj/vt/zkocc2')
       fd = tempfile.NamedTemporaryFile(dir=environment.tmproot, delete=False)