def test_service_switch(self):
     """tests the service switch from disable -> enable -> disable"""
     self._test_service_disabled()
     self._test_service_enabled()
     # 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)

 def test_service_switch(self):
   """tests the service switch from disable -> enable -> disable"""
   self._test_service_disabled()
   self._test_service_enabled()
   # 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, 'replica', 30)

    def test_service_switch(self):
        """tests the service switch from disable -> enable -> disable."""
        # make the replica spare
        utils.run_vtctl(["ChangeSlaveType", replica_tablet.tablet_alias, "spare"])
        utils.wait_for_tablet_type(replica_tablet.tablet_alias, "spare")

        # Check UpdateStreamState is disabled.
        v = utils.get_vars(replica_tablet.port)
        if v["UpdateStreamState"] != "Disabled":
            self.fail("Update stream service should be 'Disabled' but is '%s'" % v["UpdateStreamState"])

        # Make sure we can't start a new request.
        start_position = _get_repl_current_position()
        replica_conn = self._get_replica_stream_conn()
        try:
            for event in replica_conn.stream_update(
                "test_keyspace", "0", topodata_pb2.REPLICA, position=start_position
            ):
                self.assertFail("got event: %s" % str(event))
            self.assertFail("stream_update terminated with no exception")
        except dbexceptions.DatabaseError as e:
            self.assertIn("operation not allowed in state NOT_SERVING", str(e))

        # Go back to replica.
        utils.run_vtctl(["ChangeSlaveType", replica_tablet.tablet_alias, "replica"])
        utils.wait_for_tablet_type(replica_tablet.tablet_alias, "replica")

        # Check UpdateStreamState is enabled.
        v = utils.get_vars(replica_tablet.port)
        if v["UpdateStreamState"] != "Enabled":
            self.fail("Update stream service should be 'Enabled' but is '%s'" % v["UpdateStreamState"])

 def _terminated_restore(t):
   for e in utils.vtctld_connection.execute_vtctl_command(
       ['RestoreFromBackup', t.tablet_alias]):
     logging.info('%s', e.value)
     if 'shutdown mysqld' in e.value:
       break
   logging.info('waiting for restore to finish')
   utils.wait_for_tablet_type(t.tablet_alias, 'replica', timeout=30)

  def test_health_check_drained_state_does_not_shutdown_query_service(self):
    # This test is similar to test_health_check, but has the following
    # differences:
    # - the second tablet is an 'rdonly' and not a 'replica'
    # - the second tablet will be set to 'drained' and we expect that
    #   the query service won't be shutdown

    # Setup master and rdonly tablets.
    tablet_62344.init_tablet('replica', 'test_keyspace', '0')

    for t in tablet_62344, tablet_62044:
      t.create_db('vt_test_keyspace')

    # Note we only have a master and a rdonly. So we can't enable
    # semi-sync in this case, as the rdonly slaves don't semi-sync ack.
    tablet_62344.start_vttablet(wait_for_state=None, enable_semi_sync=False)
    tablet_62044.start_vttablet(wait_for_state=None,
                                init_tablet_type='rdonly',
                                init_keyspace='test_keyspace',
                                init_shard='0',
                                enable_semi_sync=False)

    tablet_62344.wait_for_vttablet_state('NOT_SERVING')
    tablet_62044.wait_for_vttablet_state('NOT_SERVING')
    self.check_healthz(tablet_62044, False)

    # Enable replication.
    utils.run_vtctl(['InitShardMaster', '-force', 'test_keyspace/0',
                     tablet_62344.tablet_alias])

    # Trigger healthcheck to save time waiting for the next interval.
    utils.run_vtctl(['RunHealthCheck', tablet_62044.tablet_alias])
    tablet_62044.wait_for_vttablet_state('SERVING')
    self.check_healthz(tablet_62044, True)

    # Change from rdonly to drained and stop replication. (These
    # actions are similar to the SplitClone vtworker command
    # implementation.)  The tablet will stay healthy, and the
    # query service is still running.
    utils.run_vtctl(['ChangeSlaveType', tablet_62044.tablet_alias, 'drained'])
    utils.run_vtctl(['StopSlave', tablet_62044.tablet_alias])
    # Trigger healthcheck explicitly to avoid waiting for the next interval.
    utils.run_vtctl(['RunHealthCheck', tablet_62044.tablet_alias])
    utils.wait_for_tablet_type(tablet_62044.tablet_alias, 'drained')
    self.check_healthz(tablet_62044, True)
    # Query service is still running.
    tablet_62044.wait_for_vttablet_state('SERVING')

    # Restart replication. Tablet will become healthy again.
    utils.run_vtctl(['ChangeSlaveType', tablet_62044.tablet_alias, 'rdonly'])
    utils.run_vtctl(['StartSlave', tablet_62044.tablet_alias])
    utils.run_vtctl(['RunHealthCheck', tablet_62044.tablet_alias])
    self.check_healthz(tablet_62044, True)

    # kill the tablets
    tablet.kill_tablets([tablet_62344, tablet_62044])

def setup_sharded_keyspace():
  utils.run_vtctl(['CreateKeyspace', SHARDED_KEYSPACE])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', SHARDED_KEYSPACE,
                   'keyspace_id', 'uint64'])

  shard_0_master.init_tablet(
      'master',
      keyspace=SHARDED_KEYSPACE,
      shard='-80',
      tablet_index=0)
  shard_0_replica.init_tablet(
      'replica',
      keyspace=SHARDED_KEYSPACE,
      shard='-80',
      tablet_index=1)
  shard_1_master.init_tablet(
      'master',
      keyspace=SHARDED_KEYSPACE,
      shard='80-',
      tablet_index=0)
  shard_1_replica.init_tablet(
      'replica',
      keyspace=SHARDED_KEYSPACE,
      shard='80-',
      tablet_index=1)

  utils.run_vtctl(['RebuildKeyspaceGraph', SHARDED_KEYSPACE], auto_log=True)

  for t in [shard_0_master, shard_0_replica, shard_1_master, shard_1_replica]:
    t.create_db('vt_test_keyspace_sharded')
    t.mquery(shard_0_master.dbname, create_vt_insert_test)
    t.start_vttablet(wait_for_state=None)

  for t in [shard_0_master, shard_1_master]:
    t.wait_for_vttablet_state('SERVING')
  for t in [shard_0_replica, shard_1_replica]:
    t.wait_for_vttablet_state('NOT_SERVING')

  utils.run_vtctl(['InitShardMaster', '-force', '%s/-80' % SHARDED_KEYSPACE,
                   shard_0_master.tablet_alias], auto_log=True)
  utils.run_vtctl(['InitShardMaster', '-force', '%s/80-' % SHARDED_KEYSPACE,
                   shard_1_master.tablet_alias], auto_log=True)

  for t in [shard_0_replica, shard_1_replica]:
    utils.wait_for_tablet_type(t.tablet_alias, 'replica')
  for t in [shard_0_master, shard_0_replica, shard_1_master, shard_1_replica]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(
      ['RebuildKeyspaceGraph', SHARDED_KEYSPACE], auto_log=True)

  utils.check_srv_keyspace('test_nj', SHARDED_KEYSPACE,
                           'Partitions(master): -80 80-\n'
                           'Partitions(rdonly): -80 80-\n'
                           'Partitions(replica): -80 80-\n')

    def _test_service_enabled(self):
        start_position = _get_repl_current_position()
        logging.debug("_test_service_enabled starting @ %s", start_position)
        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, topodata_pb2.REPLICA, 30)
        thd = threading.Thread(target=self.perform_writes, name="write_thd", args=(100,))
        thd.daemon = True
        thd.start()
        replica_conn = self._get_replica_stream_conn()

        try:
            for stream_event in replica_conn.stream_update(start_position):
                if stream_event.category == update_stream.StreamEvent.DML:
                    logging.debug("Test Service Enabled: Pass")
                    break
        except Exception as e:
            self.fail("Exception in getting stream from replica: %s\n Traceback %s" % (str(e), traceback.format_exc()))
        thd.join(timeout=30)
        replica_conn.close()

        v = utils.get_vars(replica_tablet.port)
        if v["UpdateStreamState"] != "Enabled":
            self.fail("Update stream service should be 'Enabled' but is '%s'" % v["UpdateStreamState"])
        self.assertIn("SE_DML", v["UpdateStreamEvents"])
        self.assertIn("SE_POS", v["UpdateStreamEvents"])

        logging.debug("Testing enable -> disable switch starting @ %s", start_position)
        replica_conn = self._get_replica_stream_conn()
        first = True
        txn_count = 0
        try:
            for stream_event in replica_conn.stream_update(start_position):
                if first:
                    utils.run_vtctl(["ChangeSlaveType", replica_tablet.tablet_alias, "spare"])
                    utils.wait_for_tablet_type(replica_tablet.tablet_alias, topodata_pb2.SPARE, 30)
                    first = False
                else:
                    if stream_event.category == update_stream.StreamEvent.POS:
                        txn_count += 1
                # FIXME(alainjobart) gasp, the test fails but we don't assert?
                logging.debug("Test Service Switch: FAIL")
                replica_conn.close()
                return
        except dbexceptions.DatabaseError as e:
            self.assertEqual(
                "Fatal Service Error: Disconnecting because the Update Stream " "service has been disabled", str(e)
            )
        except Exception as e:
            logging.error("Exception: %s", str(e))
            logging.error("Traceback: %s", traceback.format_exc())
            self.fail("Update stream returned error '%s'" % str(e))
        logging.debug("Streamed %d transactions before exiting", txn_count)
        replica_conn.close()

 def _test_service_disabled(self):
     start_position = _get_repl_current_position()
     logging.debug("_test_service_disabled starting @ %s", start_position)
     self._exec_vt_txn(self._populate_vt_insert_test)
     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"])
     logging.debug("dialing replica update stream service")
     replica_conn = self._get_replica_stream_conn()
     try:
         for stream_event in replica_conn.stream_update(start_position):
             break
     except Exception, e:
         logging.debug(str(e))
         self.assertIn("update stream service is not enabled", str(e))

 def _test_service_disabled(self):
   start_position = _get_repl_current_position()
   logging.debug('_test_service_disabled starting @ %s', start_position)
   self._exec_vt_txn(self._populate_vt_insert_test)
   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, 'spare')
   logging.debug('dialing replica update stream service')
   replica_conn = self._get_replica_stream_conn()
   try:
     for stream_event in replica_conn.stream_update(start_position):
       break
   except Exception, e:
     logging.debug(str(e))
     self.assertIn('update stream service is not enabled', str(e))

    def test_health_check_worker_state_does_not_shutdown_query_service(self):
        # This test is similar to test_health_check, but has the following
        # differences:
        # - the second tablet is an 'rdonly' and not a 'replica'
        # - the second tablet will be set to 'worker' and we expect that
        #   the query service won't be shutdown

        # Setup master and rdonly tablets.
        tablet_62344.init_tablet("master", "test_keyspace", "0")

        for t in tablet_62344, tablet_62044:
            t.create_db("vt_test_keyspace")

        tablet_62344.start_vttablet(wait_for_state=None, target_tablet_type="replica")
        tablet_62044.start_vttablet(
            wait_for_state=None, target_tablet_type="rdonly", init_keyspace="test_keyspace", init_shard="0"
        )

        tablet_62344.wait_for_vttablet_state("SERVING")
        tablet_62044.wait_for_vttablet_state("NOT_SERVING")
        self.check_healthz(tablet_62044, False)

        # Enable replication.
        utils.run_vtctl(["InitShardMaster", "test_keyspace/0", tablet_62344.tablet_alias])
        # Trigger healthcheck to save time waiting for the next interval.
        utils.run_vtctl(["RunHealthCheck", tablet_62044.tablet_alias, "rdonly"])
        utils.wait_for_tablet_type(tablet_62044.tablet_alias, "rdonly")
        self.check_healthz(tablet_62044, True)
        tablet_62044.wait_for_vttablet_state("SERVING")

        # Change from rdonly to worker and stop replication. (These
        # actions are similar to the SplitClone vtworker command
        # implementation.)  The tablet will become unhealthy, but the
        # query service is still running.
        utils.run_vtctl(["ChangeSlaveType", tablet_62044.tablet_alias, "worker"])
        utils.run_vtctl(["StopSlave", tablet_62044.tablet_alias])
        # Trigger healthcheck explicitly to avoid waiting for the next interval.
        utils.run_vtctl(["RunHealthCheck", tablet_62044.tablet_alias, "rdonly"])
        utils.wait_for_tablet_type(tablet_62044.tablet_alias, "worker")
        self.check_healthz(tablet_62044, False)
        # Make sure that replication got disabled.
        self.assertIn(
            ">unhealthy: replication_reporter: " "Replication is not running</span></div>", tablet_62044.get_status()
        )
        # Query service is still running.
        tablet_62044.wait_for_vttablet_state("SERVING")

        # Restart replication. Tablet will become healthy again.
        utils.run_vtctl(["ChangeSlaveType", tablet_62044.tablet_alias, "spare"])
        utils.wait_for_tablet_type(tablet_62044.tablet_alias, "spare")
        utils.run_vtctl(["StartSlave", tablet_62044.tablet_alias])
        utils.run_vtctl(["RunHealthCheck", tablet_62044.tablet_alias, "rdonly"])
        utils.wait_for_tablet_type(tablet_62044.tablet_alias, "rdonly")
        self.check_healthz(tablet_62044, True)
        tablet_62044.wait_for_vttablet_state("SERVING")

        # kill the tablets
        tablet.kill_tablets([tablet_62344, tablet_62044])

def setup_tablets():
  # Start up a master mysql and vttablet
  logging.debug('Setting up tablets')
  utils.run_vtctl(['CreateKeyspace', KEYSPACE_NAME])
  utils.run_vtctl(['SetKeyspaceShardingInfo', '-force', KEYSPACE_NAME,
                   'keyspace_id', 'uint64'])
  shard_0_master.init_tablet(
      'master',
      keyspace=KEYSPACE_NAME,
      shard='0',
      tablet_index=0)
  shard_0_replica1.init_tablet(
      'replica',
      keyspace=KEYSPACE_NAME,
      shard='0',
      tablet_index=1)

  utils.run_vtctl(['RebuildKeyspaceGraph', KEYSPACE_NAME], auto_log=True)

  for t in [shard_0_master, shard_0_replica1]:
    t.create_db('vt_test_keyspace')
    for create_table in create_tables:
      t.mquery(shard_0_master.dbname, create_table)
    t.start_vttablet(wait_for_state=None, target_tablet_type='replica')

  for t in [shard_0_master]:
    t.wait_for_vttablet_state('SERVING')
  for t in [shard_0_replica1]:
    t.wait_for_vttablet_state('NOT_SERVING')

  utils.run_vtctl(['InitShardMaster', KEYSPACE_NAME+'/0',
                   shard_0_master.tablet_alias], auto_log=True)

  for t in [shard_0_replica1]:
    utils.wait_for_tablet_type(t.tablet_alias, 'replica')

  for t in [shard_0_master, shard_0_replica1]:
    t.wait_for_vttablet_state('SERVING')

  utils.run_vtctl(
      ['RebuildKeyspaceGraph', KEYSPACE_NAME], auto_log=True)

  utils.check_srv_keyspace(
      'test_nj', KEYSPACE_NAME,
      'Partitions(master): -\n'
      'Partitions(rdonly): -\n'
      'Partitions(replica): -\n')

def setUpModule():
  try:
    environment.topo_server().setup()
    setup_procs = [master_tablet.init_mysql(),
                   replica_tablet.init_mysql()]
    utils.wait_procs(setup_procs)

    # start a vtctld so the vtctl insert commands are just RPCs, not forks.
    utils.Vtctld().start()

    # Start up a master mysql and vttablet
    logging.debug('Setting up tablets')
    utils.run_vtctl(['CreateKeyspace', 'test_keyspace'])
    master_tablet.init_tablet('master', 'test_keyspace', '0', tablet_index=0)
    replica_tablet.init_tablet('replica', 'test_keyspace', '0', tablet_index=1)
    utils.run_vtctl(['RebuildKeyspaceGraph', 'test_keyspace'], auto_log=True)
    utils.validate_topology()
    master_tablet.create_db('vt_test_keyspace')
    replica_tablet.create_db('vt_test_keyspace')

    master_tablet.start_vttablet(wait_for_state=None)
    replica_tablet.start_vttablet(wait_for_state=None)
    master_tablet.wait_for_vttablet_state('SERVING')
    replica_tablet.wait_for_vttablet_state('NOT_SERVING')
    utils.run_vtctl(['InitShardMaster', 'test_keyspace/0',
                     master_tablet.tablet_alias], auto_log=True)

    utils.wait_for_tablet_type(replica_tablet.tablet_alias, 'replica')
    master_tablet.wait_for_vttablet_state('SERVING')
    replica_tablet.wait_for_vttablet_state('SERVING')

    master_tablet.mquery('vt_test_keyspace', _create_vt_a)
    master_tablet.mquery('vt_test_keyspace', _create_vt_b)

    utils.run_vtctl(['ReloadSchema', master_tablet.tablet_alias])
    utils.run_vtctl(['ReloadSchema', replica_tablet.tablet_alias])
    utils.run_vtctl(['RebuildVSchemaGraph'])

    utils.VtGate().start(tablets=[master_tablet, replica_tablet])
    utils.vtgate.wait_for_endpoints('test_keyspace.0.master', 1)
    utils.vtgate.wait_for_endpoints('test_keyspace.0.replica', 1)

  except:
    tearDownModule()
    raise

 def _test_service_disabled(self):
     start_position = _get_repl_current_position()
     logging.debug("_test_service_disabled starting @ %s", start_position)
     self._exec_vt_txn(self._populate_vt_insert_test)
     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, "spare")
     replica_conn = self._get_replica_stream_conn()
     logging.debug("dialing replica update stream service")
     replica_conn.dial()
     try:
         data = replica_conn.stream_start(start_position)
     except Exception, e:
         logging.debug(str(e))
         if str(e) == "update stream service is not enabled":
             logging.debug("Test Service Disabled: Pass")
         else:
             self.fail("Test Service Disabled: Fail - did not throw the correct exception")

    def _test_service_enabled(self):
        start_position = _get_repl_current_position()
        logging.debug("_test_service_enabled starting @ %s", start_position)
        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)
        thd = threading.Thread(target=self.perform_writes, name="write_thd", args=(100,))
        thd.daemon = True
        thd.start()
        replica_conn = self._get_replica_stream_conn()

        try:
            for stream_event in replica_conn.stream_update(start_position):
                if stream_event.category == update_stream.StreamEvent.DML:
                    logging.debug("Test Service Enabled: Pass")
                    break
        except Exception, e:
            self.fail("Exception in getting stream from replica: %s\n Traceback %s" % (str(e), traceback.print_exc()))

    def _test_service_disabled(self):
        start_position = _get_repl_current_position()
        logging.debug("_test_service_disabled starting @ %s", start_position)
        self._exec_vt_txn(self._populate_vt_insert_test)
        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, topodata_pb2.SPARE)
        logging.debug("dialing replica update stream service")
        replica_conn = self._get_replica_stream_conn()
        try:
            for _ in replica_conn.stream_update(start_position):
                break
        except Exception as e:
            self.assertIn("update stream service is not enabled", str(e))
        replica_conn.close()

        v = utils.get_vars(replica_tablet.port)
        if v["UpdateStreamState"] != "Disabled":
            self.fail("Update stream service should be 'Disabled' but is '%s'" % v["UpdateStreamState"])

  def test_service_switch(self):
    """tests the service switch from disable -> enable -> disable."""
    # make the replica spare
    utils.run_vtctl(['ChangeSlaveType', replica_tablet.tablet_alias, 'spare'])
    utils.wait_for_tablet_type(replica_tablet.tablet_alias, 'spare')

    # Check UpdateStreamState is disabled.
    v = utils.get_vars(replica_tablet.port)
    if v['UpdateStreamState'] != 'Disabled':
      self.fail("Update stream service should be 'Disabled' but is '%s'" %
                v['UpdateStreamState'])

    start_position = _get_repl_current_position()

    # Make sure we can't start a new request to vttablet directly.
    _, stderr = utils.run_vtctl(['VtTabletUpdateStream',
                                 '-position', start_position,
                                 replica_tablet.tablet_alias],
                                expect_fail=True)
    self.assertIn('operation not allowed in state NOT_SERVING', stderr)

    # Make sure we can't start a new request through vtgate.
    replica_conn = self._get_vtgate_stream_conn()
    try:
      for event, resume_timestamp in replica_conn.update_stream(
          'test_keyspace', topodata_pb2.REPLICA,
          event=query_pb2.EventToken(shard='0', position=start_position),
          shard='0'):
        self.assertFail('got event(%d): %s' % (resume_timestamp, str(event)))
      self.assertFail('update_stream terminated with no exception')
    except dbexceptions.DatabaseError as e:
      self.assertIn(vtgate_gateway_flavor().no_tablet_found_message(), str(e))

    # Go back to replica.
    utils.run_vtctl(
        ['ChangeSlaveType', replica_tablet.tablet_alias, 'replica'])
    utils.wait_for_tablet_type(replica_tablet.tablet_alias, 'replica')

    # Check UpdateStreamState is enabled.
    v = utils.get_vars(replica_tablet.port)
    if v['UpdateStreamState'] != 'Enabled':
      self.fail("Update stream service should be 'Enabled' but is '%s'" %
                v['UpdateStreamState'])

  def test_update_stream_interrupt(self):
    """Checks that a running query is terminated on going non-serving."""
    # Make sure the replica is replica type.
    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, 'replica', 30)

    # Save current position, insert some data.
    start_position = _get_repl_current_position()
    logging.debug('test_update_stream_interrupt starting @ %s', start_position)
    self._exec_vt_txn(self._populate_vt_a(1))
    self._exec_vt_txn(['delete from vt_a'])

    # Start an Update Stream from the slave. When we get the data, go to spare.
    # That should interrupt the streaming RPC.
    replica_conn = self._get_vtgate_stream_conn()
    first = True
    txn_count = 0
    try:
      for event, resume_timestamp in replica_conn.update_stream(
          'test_keyspace', topodata_pb2.REPLICA,
          event=query_pb2.EventToken(shard='0', position=start_position),
          shard='0'):
        logging.debug('test_update_stream_interrupt got event(%d): %s',
                      resume_timestamp, event)
        if first:
          utils.run_vtctl(
              ['ChangeSlaveType', replica_tablet.tablet_alias, 'spare'])
          utils.wait_for_tablet_type(replica_tablet.tablet_alias, 'spare', 30)
          first = False
        else:
          if event.event_token.position:
            txn_count += 1

      self.assertFail('update_stream terminated with no exception')
    except dbexceptions.DatabaseError as e:
      self.assertIn('context canceled', str(e))
    self.assertFalse(first)

    logging.debug('Streamed %d transactions before exiting', txn_count)
    replica_conn.close()

  def _test_service_enabled(self):
    start_position = _get_repl_current_position()
    logging.debug('_test_service_enabled starting @ %s', start_position)
    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, 'replica', 30)
    thd = threading.Thread(target=self.perform_writes, name='write_thd',
                           args=(100,))
    thd.daemon = True
    thd.start()
    replica_conn = self._get_replica_stream_conn()

    try:
      for stream_event in replica_conn.stream_update(start_position):
        if stream_event.category == update_stream.StreamEvent.DML:
          logging.debug('Test Service Enabled: Pass')
          break
    except Exception, e:
      self.fail('Exception in getting stream from replica: %s\n Traceback %s' %
                (str(e), traceback.print_exc()))

  def _test_service_disabled(self):
    start_position = _get_repl_current_position()
    logging.debug('_test_service_disabled starting @ %s', start_position)
    self._exec_vt_txn(self._populate_vt_insert_test)
    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, 'spare')
    logging.debug('dialing replica update stream service')
    replica_conn = self._get_replica_stream_conn()
    try:
      for _ in replica_conn.stream_update(start_position):
        break
    except dbexceptions.DatabaseError as e:
      self.assertIn('update stream service is not enabled', str(e))
    replica_conn.close()

    v = utils.get_vars(replica_tablet.port)
    if v['UpdateStreamState'] != 'Disabled':
      self.fail("Update stream service should be 'Disabled' but is '%s'" %
                v['UpdateStreamState'])

    def _test_service_enabled(self):
        start_position = _get_repl_current_position()
        logging.debug("_test_service_enabled starting @ %s", start_position)
        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, "replica", 30)
        thd = threading.Thread(target=self.perform_writes, name="write_thd", args=(100,))
        thd.daemon = True
        thd.start()
        replica_conn = self._get_replica_stream_conn()
        replica_conn.dial()

        try:
            data = replica_conn.stream_start(start_position)
            for i in xrange(10):
                data = replica_conn.stream_next()
                if data["Category"] == "DML" and utils.options.verbose == 2:
                    logging.debug("Test Service Enabled: Pass")
                    break
        except Exception, e:
            self.fail("Exception in getting stream from replica: %s\n Traceback %s" % (str(e), traceback.print_exc()))