Make topo server flavor pluggable in integration tests.

2014-11-14 16:37:05 -08:00 · 2014-11-14 16:37:05 -08:00 · 3ee8a59cf1
--- a/test/binlog.py
+++ b/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()

--- a/test/clone.py
+++ b/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()
--- a/test/environment.py
+++ b/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']
--- a/test/initial_sharding.py
+++ b/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()

--- a/test/java_vtgate_test_helper.py
+++ b/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:
--- a/test/keyspace_test.py
+++ b/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()

--- a/test/mysqlctl.py
+++ b/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()

--- a/test/primecache.py
+++ b/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()

--- a/test/queryservice_test.py
+++ b/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:
--- a/test/queryservice_tests/test_env.py
+++ b/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()
--- a/test/reparent.py
+++ b/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',
--- a/test/resharding.py
+++ b/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()

--- a/test/rowcache_invalidator.py
+++ b/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()
--- a/test/schema.py
+++ b/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()

--- a/test/sharded.py
+++ b/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])
--- a/test/tablet.py
+++ b/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)
--- a/test/tabletmanager.py
+++ b/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

--- a/test/topo_flavor/init.py
+++ b/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.
--- a/test/topo_flavor/server.py
+++ b/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
--- a/test/topo_flavor/zookeeper.py
+++ b/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()
--- a/test/update_stream.py
+++ b/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()
--- a/test/utils.py
+++ b/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)
--- a/test/vertical_split.py
+++ b/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()

--- a/test/vtctld_test.py
+++ b/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
--- a/test/vtdb_test.py
+++ b/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()
--- a/test/vtgatev2_test.py
+++ b/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()
--- a/test/zkocc_test.py
+++ b/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()
-    if environment.topo_server_implementation == 'zookeeper':
+    environment.topo_server().wipe()
+    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)