import org.apache.zookeeper.*;
import org.apache.zookeeper.KeeperException.*;
import org.apache.zookeeper.data.*;
import java.io.IOException;
import java.util.*;

class zktest {
    static Watcher nullWatcher = new Watcher() {
        public void process(WatchedEvent event) {
        }
    };

    public static void kill(ZooKeeper victim) throws IOException, KeeperException, InterruptedException {
        ZooKeeper newHandle = new ZooKeeper("127.0.0.1:2181", 100000, nullWatcher,
                                            victim.getSessionId(), victim.getSessionPasswd());
        newHandle.close();
    }

    public static void main(String[] args) throws IOException, KeeperException, InterruptedException {
        System.out.println("Preparing test");

        ZooKeeper zk = new ZooKeeper("127.0.0.1:2181", 100000, nullWatcher);
        try {
            zk.create("/test", null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        } catch (NodeExistsException e) {
            // pass
        }

        System.out.println("Making sure this client is fully operational");
        zk.setData("/test", new byte[0], -1);

        System.out.println("Killing session");
        kill(zk);

        System.out.println("Trying a setData on the dead session");
        while (true) {
            try {
                zk.setData("/test", new byte[0], -1);
            } catch (ConnectionLossException e) {
                System.out.println("Connection lost, retrying");
                continue;
            } catch (SessionExpiredException e) {
                System.out.println("Session expired");
                System.exit(0);
            }
            System.out.println("ERROR: Write succeeded, but it shouldn't have");
            System.exit(1);
        }
    }
}
