com.comcast.viper.flume2storm.zookeeper.ZkTestUtils.java Source code

Introduction

Here is the source code for com.comcast.viper.flume2storm.zookeeper.ZkTestUtils.java

Source

/**
 * Copyright 2014 Comcast Cable Communications Management, LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.comcast.viper.flume2storm.zookeeper;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.Socket;

import org.apache.commons.lang.StringUtils;

import com.comcast.viper.flume2storm.utility.test.TestCondition;
import com.comcast.viper.flume2storm.utility.test.TestUtils;
import com.google.common.base.Preconditions;

/**
 * 
 */
public class ZkTestUtils {
    public static final int TEST_TIMEOUT = 10000;
    public static final int ZK_OP_TIMEOUT = 3000;
    public static final int A_LITTLE = 1500;

    public static final String HOST = "127.0.0.1";
    public static final int PORT = TestUtils.getAvailablePort();
    public static final String ZK_OK_CMD = "ruok";

    // This method is copied from Flume OG's code
    /**
     * Utility function to send a command to the internal server. ZK servers
     * accepts 4 byte command strings to test their liveness.
     */
    protected static String send4LetterWord(final String host, final int port, final String cmd) {
        Preconditions.checkArgument(cmd.length() == 4);
        try {
            final Socket sock = new Socket(host, port);
            OutputStream outstream = null;
            BufferedReader reader = null;
            try {
                outstream = sock.getOutputStream();
                outstream.write(cmd.getBytes());
                outstream.flush();

                reader = new BufferedReader(new InputStreamReader(sock.getInputStream()));
                final StringBuilder sb = new StringBuilder();
                String line;
                while ((line = reader.readLine()) != null) {
                    sb.append(line + "\n");
                }
                return sb.toString();
            } finally {
                if (outstream != null) {
                    outstream.close();
                }
                sock.close();
                if (reader != null) {
                    reader.close();
                }
            }
        } catch (final Exception e) {
            System.out.println(e.getMessage());
            return StringUtils.EMPTY;
        }
    }

    public static final TestCondition ZK_IS_ON = new TestCondition() {
        @Override
        public boolean evaluate() {
            return !send4LetterWord(HOST, PORT, ZK_OK_CMD).isEmpty();
        }
    };

    public static final TestCondition ZK_IS_OFF = new TestCondition() {
        @Override
        public boolean evaluate() {
            return send4LetterWord(HOST, PORT, ZK_OK_CMD).isEmpty();
        }
    };
}