Kazoo test for create()

gistfile1.py

import kazoo.client
import kazoo.exceptions
from kazoo.protocol.paths import join as zjoin
import threading
import time


# =====
class TimeOf:
    def __init__(self, name, print_result=True):
        self._name = name
        self._print_result = print_result
        self._before = None
        self._result = None

    def get_result(self):
        return self._result

    def __enter__(self):
        if self._print_result:
            print("--> Entering to the region '{}'...".format(self._name))
        self._before = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._result = time.time() - self._before
        if self._print_result:
            print("<-- Region '{}' completed, execution time: {:0.5f} sec\n".format(self._name, self._result))

def zoo_connect():
    zk = kazoo.client.KazooClient()
    zk.start()
    return zk

def action(zoo_connect, name, method_name, root_path, offset, limit):
    zk = zoo_connect()
    with TimeOf(name, False) as time_of:
        total = 0
        for count in range(offset, limit):
            with TimeOf(name + str(count), False) as time_of_op:
                method = getattr(zk, method_name)
                method("{}/node-{}".format(root_path, count))
            total += time_of_op.get_result()
    op_time = (total / (limit - offset))
    print("-- {}\tregion:  {:0.5f} sec".format(name, time_of.get_result()))
    print("-- {}\top-time: {:0.5f} sec".format(name, op_time))
    print("-- {}\trps:     {:0.5f}".format(name, 1 / op_time))

# =====
def main():
    root_path = "/perf-test"
    threads_number = 20
    nodes_per_thread = 500

    with TimeOf("connect"):
        zk = zoo_connect()

    with TimeOf("cleannup"):
        zk.delete(root_path, recursive=True)

    with TimeOf("init"):
        zk.create(root_path)

    import multiprocessing
    for method_name in ("create",):# "delete"):
        with TimeOf("threads_" + method_name):
            threads = [
                #threading.Thread(
                multiprocessing.Process(
                    target=action,
                    args=(
                        zoo_connect,
                        "thread_{}_{}".format(method_name, count),
                        method_name,
                        root_path,
                        count * nodes_per_thread,
                        count * nodes_per_thread + nodes_per_thread,
                    ),
                )
                for count in range(threads_number)
            ]
            for thread in threads:
                thread.start()
            for thread in threads:
                thread.join()


if __name__ == "__main__":
    main()