/* Copyright 2013 Endgame, Inc.
 * 
 * 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 org.elasticsearch.river.kafka;

import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.aryEq;
import static org.easymock.EasyMock.createMock;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.easymock.EasyMock.expectLastCall;
import static org.easymock.EasyMock.replay;
import static org.easymock.EasyMock.verify;

import java.util.Collections;

import junit.framework.TestCase;
import kafka.api.FetchRequest;
import kafka.javaapi.consumer.SimpleConsumer;
import kafka.javaapi.message.ByteBufferMessageSet;

import org.apache.zookeeper.data.Stat;

import com.netflix.curator.framework.CuratorFramework;
import com.netflix.curator.framework.api.ExistsBuilder;
import com.netflix.curator.framework.api.GetDataBuilder;
import com.netflix.curator.framework.api.SetDataBuilder;

public class KafkaClientTest extends TestCase {
	
	public ExistsBuilder mockExistsBuilder(String path, boolean exists) throws Exception
	{
		ExistsBuilder b = createMock(ExistsBuilder.class);
		if(exists)
		{
			expect(b.forPath(eq(path))).andReturn(new Stat());
		}
		else
		{
			expect(b.forPath(eq(path))).andReturn(null);
		}
		replay(b);		
		return b;
	}
	
	public GetDataBuilder mockGetDataBuilder(String path, byte[] data) throws Exception
	{
		GetDataBuilder b = createMock(GetDataBuilder.class);
		expect(b.forPath(eq(path))).andReturn(data);
		replay(b);
		return b;
	}
	
	public SetDataBuilder mockSetDataBuilder(String path, byte[] data) throws Exception
	{
		SetDataBuilder b = createMock(SetDataBuilder.class);
		expect(b.forPath(eq(path), aryEq(data))).andReturn(new Stat());
		replay(b);
		return b;
	}
	
	CuratorFramework mockCurator;
	SimpleConsumer mockConsumer;
	KafkaClient client;
	
	@Override
	protected void setUp() throws Exception {
		super.setUp();
		mockCurator = createMock(CuratorFramework.class);
		mockConsumer = createMock(SimpleConsumer.class);
		
		final CuratorFramework cur = mockCurator;
		final SimpleConsumer con = mockConsumer;
		client = new KafkaClient("zookeeper", "broker", 9092){
			void connect(String zk, String broker, int port) 
			{
				this.curator = cur;
				this.consumer = con;
			};
		};
	}
	
	@Override
	protected void tearDown() throws Exception {
		super.tearDown();
		mockCurator = null;
		mockConsumer = null;
		client = null;
	}

	public void testConstructor()
	{
		replay(mockConsumer, mockCurator);
		assertEquals(client.brokerURL, "broker:9092");
	}
	
	public void testGetNewestOffset() throws Exception
	{
		long[] mylong = new long[3];
		mylong[0] = 10;
		mylong[1] = 20;
		mylong[2] = 30;
		
		expect(mockConsumer.getOffsetsBefore("topic", 1, -1, 1)).andReturn(mylong);	
		replay(mockConsumer, mockCurator);		
		long answer = client.getNewestOffset("topic", 1);
		assertEquals(mylong[0], answer);
		verify(mockConsumer, mockCurator);
	}

	public void testGetOldestOffset() throws Exception
	{
		long[] mylong = new long[3];
		mylong[0] = 10;
		mylong[1] = 20;
		mylong[2] = 30;
		
		expect(mockConsumer.getOffsetsBefore("topic", 1, -2, 1)).andReturn(mylong);	
		replay(mockConsumer, mockCurator);		
		long answer = client.getOldestOffset("topic", 1);
		assertEquals(mylong[0], answer);
		verify(mockConsumer, mockCurator);
	}
	
	public void testGet() throws Exception
	{
		expect(mockCurator.checkExists()).andReturn(mockExistsBuilder("/some/path", false));
		expect(mockCurator.checkExists()).andReturn(mockExistsBuilder("/some/existing/path", true));
		expect(mockCurator.getData()).andReturn(mockGetDataBuilder("/some/existing/path", "this is data".getBytes()));
		
		replay(mockConsumer, mockCurator);
		assertNull(client.get("/some/path"));
		assertEquals(client.get("/some/existing/path"), "this is data");
		verify(mockConsumer, mockCurator);
	}
	
	public void testSave() throws Exception
	{
		expect(mockCurator.checkExists()).andReturn(mockExistsBuilder("/some/existing/path", true));
		expect(mockCurator.setData()).andReturn(mockSetDataBuilder("/some/existing/path", "this is data".getBytes()));
		
		replay(mockConsumer, mockCurator);
		client.save("/some/existing/path", "this is data");
		verify(mockConsumer, mockCurator);
	}
	
	static class Args
	{
		String path = "";
		String data = "";
	}
	
	public void testSaveOffset()
	{
		final Args args = new Args();
		replay(mockConsumer, mockCurator);
		client = new KafkaClient("zookeeper", "broker", 9092){
			void connect(String zk, String broker, int port) 
			{
				this.curator = mockCurator;
				this.consumer = mockConsumer;
			}
			
			@Override
			public void save(String path, String data) {
				args.path = path;
				args.data = data;
			}
		};
		client.saveOffset("my_topic", 77, 4242);

		assertEquals("/es-river-kafka/offsets/broker:9092/my_topic/77", args.path);
		assertEquals("4242", args.data);
	}
	
	public void testGetOffsets()
	{
		final Args args = new Args();
		replay(mockConsumer, mockCurator);
		client = new KafkaClient("zookeeper", "broker", 9092){
			void connect(String zk, String broker, int port) 
			{
				this.curator = mockCurator;
				this.consumer = mockConsumer;
			}
		
			@Override
			public String get(String path) {
				args.path = path;
				return "100";
			}
		};
		
		assertEquals(100, client.getOffset("my_topic", 777));
		assertEquals("/es-river-kafka/offsets/broker:9092/my_topic/777", args.path);
	}
	
	public void testClose()
	{
		mockCurator.close();
		expectLastCall().asStub();
		replay(mockConsumer, mockCurator);
		
		client.close();
		verify(mockConsumer, mockCurator);
	}
	
	public void testFetch()
	{
		expect(mockConsumer.fetch(anyObject(FetchRequest.class))).andReturn(new ByteBufferMessageSet(Collections.EMPTY_LIST));	
		replay(mockConsumer, mockCurator);		
		client.fetch("my_topic", 0, 1717, 1024);
		verify(mockConsumer, mockCurator);
	}
	
}