"""Kazoo testing harnesses"""
import atexit
import logging
import os
import uuid
import unittest

from kazoo.client import KazooClient
from kazoo.exceptions import NotEmptyError
from kazoo.protocol.states import (
    KazooState
)
from kazoo.testing.common import ZookeeperCluster
from kazoo.protocol.connection import _CONNECTION_DROP, _SESSION_EXPIRED

log = logging.getLogger(__name__)

CLUSTER = None


def get_global_cluster():
    global CLUSTER
    if CLUSTER is None:
        ZK_HOME = os.environ.get("ZOOKEEPER_PATH")
        ZK_CLASSPATH = os.environ.get("ZOOKEEPER_CLASSPATH")
        ZK_PORT_OFFSET = int(os.environ.get("ZOOKEEPER_PORT_OFFSET", 20000))

        assert ZK_HOME or ZK_CLASSPATH, (
            "Either ZOOKEEPER_PATH or ZOOKEEPER_CLASSPATH environment "
            "variable must be defined.\n"
            "For deb package installations this is /usr/share/java")

        CLUSTER = ZookeeperCluster(
            install_path=ZK_HOME,
            classpath=ZK_CLASSPATH,
            port_offset=ZK_PORT_OFFSET,
        )
        atexit.register(lambda cluster: cluster.terminate(), CLUSTER)
    return CLUSTER


class KazooTestHarness(unittest.TestCase):
    """Harness for testing code that uses Kazoo

    This object can be used directly or as a mixin. It supports starting
    and stopping a complete ZooKeeper cluster locally and provides an
    API for simulating errors and expiring sessions.

    Example::

        class MyTestCase(KazooTestHarness):
            def setUp(self):
                self.setup_zookeeper()

                # additional test setup

            def tearDown(self):
                self.teardown_zookeeper()

            def test_something(self):
                something_that_needs_a_kazoo_client(self.client)

            def test_something_else(self):
                something_that_needs_zk_servers(self.servers)

    """

    def __init__(self, *args, **kw):
        super(KazooTestHarness, self).__init__(*args, **kw)
        self.client = None
        self._clients = []

    @property
    def cluster(self):
        return get_global_cluster()

    @property
    def servers(self):
        return ",".join([s.address for s in self.cluster])

    def _get_nonchroot_client(self):
        return KazooClient(self.servers)

    def _get_client(self, **kwargs):
        c = KazooClient(self.hosts, **kwargs)
        try:
            self._clients.append(c)
        except AttributeError:
            self._client = [c]
        return c

    def lose_connection(self, event_factory):
        """Force client to lose connection with server"""
        self.__break_connection(_CONNECTION_DROP, KazooState.SUSPENDED, event_factory)

    def expire_session(self, event_factory):
        """Force ZK to expire a client session"""
        self.__break_connection(_SESSION_EXPIRED, KazooState.LOST, event_factory)

    def setup_zookeeper(self, **client_options):
        """Create a ZK cluster and chrooted :class:`KazooClient`

        The cluster will only be created on the first invocation and won't be
        fully torn down until exit.
        """
        if not self.cluster[0].running:
            self.cluster.start()
        namespace = "/kazootests" + uuid.uuid4().hex
        self.hosts = self.servers + namespace

        if 'timeout' not in client_options:
            client_options['timeout'] = 0.8
        self.client = self._get_client(**client_options)
        self.client.start()
        self.client.ensure_path("/")

    def teardown_zookeeper(self):
        """Clean up any ZNodes created during the test
        """
        if not self.cluster[0].running:
            self.cluster.start()

        tries = 0
        if self.client and self.client.connected:
            while tries < 3:
                try:
                    self.client.retry(self.client.delete, '/', recursive=True)
                    break
                except NotEmptyError:
                    pass
                tries += 1
            self.client.stop()
            self.client.close()
            del self.client
        else:
            client = self._get_client()
            client.start()
            client.retry(client.delete, '/', recursive=True)
            client.stop()
            client.close()
            del client

        for client in self._clients:
            client.stop()
            del client
        self._clients = None

    def __break_connection(self, break_event, expected_state, event_factory):
        """Break ZooKeeper connection using the specified event."""

        lost = event_factory()
        safe = event_factory()

        def watch_loss(state):
            if state == expected_state:
                lost.set()
            elif lost.is_set() and state == KazooState.CONNECTED:
                safe.set()
                return True

        self.client.add_listener(watch_loss)
        self.client._call(break_event, None)

        lost.wait(5)
        if not lost.isSet():
            raise Exception("Failed to get notified of broken connection.")

        safe.wait(15)
        if not safe.isSet():
            raise Exception("Failed to see client reconnect.")

        self.client.retry(self.client.get_async, '/')


class KazooTestCase(KazooTestHarness):
    def setUp(self):
        self.setup_zookeeper()

    def tearDown(self):
        self.teardown_zookeeper()

@@ -27,21 +27,8 @@
 %22%22%22%0A
-import atexit
 %0Aimp
@@ -78,16 +78,59 @@
 ittest%0A%0A
+from kazoo import python2atexit as atexit%0A%0A
 from kaz

Use the compat layer for atexit functionality