import time
from multiprocessing import Event

from perfrunner.tests import PerfTest
from perfrunner.tests.index import IndexTest


def with_delay(method):
    def wrapper(self, *args, **kwargs):
        time.sleep(self.rebalance_settings.start_after)
        method(self, *args, **kwargs)
        time.sleep(self.rebalance_settings.stop_after)
        self.shutdown_event.set()
    return wrapper


class RebalanceTest(PerfTest):

    def __init__(self, *args, **kwargs):
        super(RebalanceTest, self).__init__(*args, **kwargs)
        self.shutdown_event = Event()
        self.rebalance_settings = self.test_config.get_rebalance_settings()

    @with_delay
    def rebalance_in(self):
        for cluster in self.cluster_spec.get_clusters():
            master = cluster[0]
            known_nodes = cluster[:self.rebalance_settings.nodes_after]
            ejected_nodes = []
            self.rest.rebalance(master, known_nodes, ejected_nodes)
            self.monitor.monitor_rebalance(master)


class StaticRebalanceTest(RebalanceTest):

    def run(self):
        self._run_load_phase()
        self._compact_bucket()

        self.reporter.start()
        self.rebalance_in()
        value = self.reporter.finish('Rebalance')
        self.reporter.post_to_sf(self, value)

        self._debug()


class StaticRebalanceWithIndexTest(RebalanceTest, IndexTest):

    def run(self):
        self._run_load_phase()
        self._compact_bucket()

        self._define_ddocs()
        self._build_index()

        self.reporter.start()
        self.rebalance_in()
        value = self.reporter.finish('Rebalance')
        self.reporter.post_to_sf(self, value)

        self._debug()

@@ -678,32 +678,151 @@
 lance_in(self):%0A
+        initial_nodes = self.test_config.get_initial_nodes()%0A        nodes_after = self.rebalance_settings.nodes_after%0A
         for clus
@@ -894,16 +894,17 @@
 ster%5B0%5D%0A
+%0A
         
@@ -911,55 +911,193 @@
     
-known_nodes = cluster%5B:self.rebalance_settings.
+for host_port in cluster%5Binitial_nodes:nodes_after%5D:%0A                host = host_port.split(':')%5B0%5D%0A                self.rest.add_node(master, host)%0A%0A            known_nodes = cluster%5B:
 node

add nodes before rebalance