Java tutorial
/** * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You 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.pinterest.terrapin.zookeeper; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import static org.mockito.Matchers.any; import static org.mockito.Matchers.anyBoolean; import static org.mockito.Matchers.anyInt; import static org.mockito.Matchers.anyListOf; import static org.mockito.Matchers.anyString; import static org.mockito.Matchers.eq; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.powermock.api.mockito.PowerMockito.verifyStatic; import com.pinterest.terrapin.thrift.generated.Options; import com.google.common.base.Function; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Sets; import com.twitter.common.zookeeper.ZooKeeperClient; import com.twitter.common.zookeeper.ZooKeeperMap; import org.apache.commons.lang3.tuple.Pair; import org.apache.zookeeper.CreateMode; import org.apache.zookeeper.KeeperException; import org.apache.zookeeper.ZooKeeper; import org.apache.zookeeper.data.ACL; import org.apache.zookeeper.data.Stat; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.ArgumentCaptor; import org.powermock.api.mockito.PowerMockito; import org.powermock.api.support.membermodification.MemberModifier; import org.powermock.core.classloader.annotations.PrepareForTest; import org.powermock.modules.junit4.PowerMockRunner; import java.net.InetSocketAddress; import java.util.Map; import java.util.Set; @RunWith(PowerMockRunner.class) public class ZooKeeperManagerTest { private static final String CLUSTER_NAME = "test_cluster"; private static final String FILE_SET = "test_fileset"; private static final FileSetInfo FILE_SET_INFO = new FileSetInfo(); private static final String FILE_SET_DIR = "/" + CLUSTER_NAME + "/filesets"; private static final String VIEWS_DIR = "/" + CLUSTER_NAME + "/views"; private static final String LOCKS_DIR = "/" + CLUSTER_NAME + "/locks"; private static final String FILE_SET_PATH = FILE_SET_DIR + "/" + FILE_SET; private static final String FILE_SET_LOCK_PATH = LOCKS_DIR + "/" + FILE_SET; private ZooKeeperManager zkManager; private ZooKeeperClient zkClient; private ZooKeeper zk; @Before public void setUp() throws Exception { zkClient = mock(ZooKeeperClient.class); zk = mock(ZooKeeper.class); zkManager = new ZooKeeperManager(zkClient, CLUSTER_NAME); when(zkClient.get()).thenReturn(zk); } @Test public void testCreateClusterPaths() throws Exception { when(zk.create(anyString(), any(byte[].class), anyListOf(ACL.class), any(CreateMode.class))).thenReturn(""); ArgumentCaptor<String> pathCaptor = ArgumentCaptor.forClass(String.class); zkManager.createClusterPaths(); verify(zk, times(3)).create(pathCaptor.capture(), any(byte[].class), anyListOf(ACL.class), any(CreateMode.class)); Set<String> allPaths = Sets.newHashSet(pathCaptor.getAllValues()); assertTrue(allPaths.contains(FILE_SET_DIR)); assertTrue(allPaths.contains(VIEWS_DIR)); assertTrue(allPaths.contains(LOCKS_DIR)); } @Test @PrepareForTest(ZooKeeperMap.class) public void testRegisterWatchAllFileSets() throws Exception { PowerMockito.mockStatic(ZooKeeperMap.class); when(ZooKeeperMap.create(any(ZooKeeperClient.class), anyString(), any(Function.class))).thenReturn(null); ArgumentCaptor<String> pathCaptor = ArgumentCaptor.forClass(String.class); zkManager.registerWatchAllFileSets(); verifyStatic(times(2)); ZooKeeperMap.create(any(ZooKeeperClient.class), pathCaptor.capture(), any(Function.class)); Set<String> allPaths = Sets.newHashSet(pathCaptor.getAllValues()); assertTrue(allPaths.contains(FILE_SET_DIR)); assertTrue(allPaths.contains(VIEWS_DIR)); } @Test public void testLockFileSet() throws Exception { when(zk.create(anyString(), any(byte[].class), anyListOf(ACL.class), any(CreateMode.class))).thenReturn(""); CreateMode mode = CreateMode.EPHEMERAL; zkManager.lockFileSet(FILE_SET, FILE_SET_INFO, mode); verify(zk).create(eq(FILE_SET_LOCK_PATH), any(byte[].class), anyListOf(ACL.class), eq(mode)); } @Test public void testUnlockFileSet() throws Exception { doNothing().when(zk).delete(anyString(), anyInt()); ArgumentCaptor<String> pathCaptor = ArgumentCaptor.forClass(String.class); ArgumentCaptor<Integer> versionCaptor = ArgumentCaptor.forClass(Integer.class); zkManager.unlockFileSet(FILE_SET); verify(zk).delete(pathCaptor.capture(), versionCaptor.capture()); assertEquals(FILE_SET_LOCK_PATH, pathCaptor.getValue()); assertEquals(new Integer(-1), versionCaptor.getValue()); } @Test public void testSetNonExistingFileSetInfo() throws Exception { when(zk.exists(anyString(), anyBoolean())).thenReturn(null); when(zk.create(anyString(), any(byte[].class), anyListOf(ACL.class), any(CreateMode.class))).thenReturn(""); zkManager.setFileSetInfo(FILE_SET, FILE_SET_INFO); verify(zk).exists(eq(FILE_SET_PATH), eq(false)); verify(zk).create(eq(FILE_SET_PATH), eq(FILE_SET_INFO.toJson()), anyListOf(ACL.class), any(CreateMode.class)); } @Test public void testSetExistingFileSetInfo() throws Exception { when(zk.exists(anyString(), anyBoolean())).thenReturn(mock(Stat.class)); when(zk.setData(anyString(), any(byte[].class), anyInt())).thenReturn(mock(Stat.class)); zkManager.setFileSetInfo(FILE_SET, FILE_SET_INFO); verify(zk).exists(eq(FILE_SET_PATH), eq(false)); verify(zk).setData(eq(FILE_SET_PATH), eq(FILE_SET_INFO.toJson()), eq(-1)); } @Test public void testDeleteFileSetInfo() throws Exception { doNothing().when(zk).delete(anyString(), anyInt()); zkManager.deleteFileSetInfo(FILE_SET); verify(zk).delete(eq(FILE_SET_PATH), eq(-1)); } @Test public void testGetViewInfo() throws Exception { PowerMockito.mockStatic(ZooKeeperMap.class); ViewInfo viewInfo = new ViewInfo(); String resource1 = "test_view"; MemberModifier.field(ZooKeeperManager.class, "viewInfoMap").set(zkManager, ImmutableMap.of(resource1, viewInfo)); assertNotNull(zkManager.getViewInfo(resource1)); assertNull(zkManager.getViewInfo("non_existing")); } @Test public void testSetExistingViewInfo() throws Exception { ViewInfo viewInfo = mock(ViewInfo.class); String resource = "test_resource"; String resourcePath = VIEWS_DIR + "/" + resource; byte[] resourceData = "hello".getBytes(); when(viewInfo.toCompressedJson()).thenReturn(resourceData); when(viewInfo.hasOnlinePartitions()).thenReturn(true); when(viewInfo.getResource()).thenReturn(resource); when(zk.exists(eq(resourcePath), anyBoolean())).thenReturn(new Stat()); when(zk.setData(eq(resourcePath), eq(resourceData), eq(-1))).thenReturn(new Stat()); zkManager.setViewInfo(viewInfo); verify(zk).setData(eq(resourcePath), eq(resourceData), eq(-1)); } @Test public void testSetNonExistingViewInfo() throws Exception { ViewInfo viewInfo = mock(ViewInfo.class); String resource = "test_resource"; String resourcePath = VIEWS_DIR + "/" + resource; byte[] resourceData = "hello".getBytes(); when(viewInfo.toCompressedJson()).thenReturn(resourceData); when(viewInfo.hasOnlinePartitions()).thenReturn(true); when(viewInfo.getResource()).thenReturn(resource); when(zk.exists(eq(resourcePath), anyBoolean())).thenReturn(null); when(zk.create(eq(resourcePath), eq(resourceData), anyListOf(ACL.class), any(CreateMode.class))) .thenReturn(null); zkManager.setViewInfo(viewInfo); verify(zk).create(eq(resourcePath), eq(resourceData), anyListOf(ACL.class), any(CreateMode.class)); } @Test public void testSetViewInfoNoOnlinePart() throws Exception { ViewInfo viewInfo = mock(ViewInfo.class); byte[] resourceData = "hello".getBytes(); when(viewInfo.toCompressedJson()).thenReturn(resourceData); when(viewInfo.hasOnlinePartitions()).thenReturn(false); zkManager.setViewInfo(viewInfo); verify(zk, times(0)).create(anyString(), any(byte[].class), anyListOf(ACL.class), any(CreateMode.class)); verify(zk, times(0)).setData(anyString(), any(byte[].class), anyInt()); } @Test public void testGetCandidateHdfsDirMap() throws Exception { FileSetInfo fileSetInfo1 = new FileSetInfo("abc", "123", 0, ImmutableList.copyOf(new FileSetInfo.ServingInfo[] {}), new Options()); FileSetInfo fileSetInfo2 = new FileSetInfo("def", "456", 0, ImmutableList.copyOf(new FileSetInfo.ServingInfo[] {}), new Options()); FileSetInfo fileSetInfo3 = new FileSetInfo("ghi", "789", 0, ImmutableList.copyOf(new FileSetInfo.ServingInfo[] {}), new Options()); MemberModifier.field(ZooKeeperManager.class, "fileSetInfoMap").set(zkManager, ImmutableMap.builder().put(fileSetInfo1.fileSetName, fileSetInfo1) .put(fileSetInfo2.fileSetName, fileSetInfo2).put(fileSetInfo3.fileSetName, fileSetInfo3) .build()); when(zk.getData(eq(FILE_SET_DIR + "/" + fileSetInfo1.fileSetName), anyBoolean(), any(Stat.class))) .thenReturn(fileSetInfo1.toJson()); when(zk.getData(eq(LOCKS_DIR + "/" + fileSetInfo1.fileSetName), anyBoolean(), any(Stat.class))) .thenReturn(fileSetInfo1.toJson()); when(zk.getData(eq(FILE_SET_DIR + "/" + fileSetInfo2.fileSetName), anyBoolean(), any(Stat.class))) .thenReturn(fileSetInfo2.toJson()); when(zk.getData(eq(LOCKS_DIR + "/" + fileSetInfo2.fileSetName), anyBoolean(), any(Stat.class))) .thenThrow(new KeeperException.NoNodeException()); when(zk.getData(eq(FILE_SET_DIR + "/" + fileSetInfo3.fileSetName), anyBoolean(), any(Stat.class))) .thenThrow(new KeeperException.NoNodeException()); when(zk.getData(eq(LOCKS_DIR + "/" + fileSetInfo3.fileSetName), anyBoolean(), any(Stat.class))) .thenReturn(fileSetInfo3.toJson()); Map<String, Pair<FileSetInfo, FileSetInfo>> results = zkManager.getCandidateHdfsDirMap(); assertEquals(3, results.size()); assertEquals(fileSetInfo1, results.get(fileSetInfo1.fileSetName).getLeft()); assertEquals(fileSetInfo1, results.get(fileSetInfo1.fileSetName).getRight()); assertEquals(fileSetInfo2, results.get(fileSetInfo2.fileSetName).getLeft()); assertNull(results.get(fileSetInfo2.fileSetName).getRight()); assertNull(results.get(fileSetInfo3.fileSetName).getLeft()); assertEquals(fileSetInfo3, results.get(fileSetInfo3.fileSetName).getRight()); } @Test public void testGetControllerLeader() throws Exception { String host = "somehost"; int port = 9090; String json = String.format("{\"id\": \"%s_%d\"}", host, port); when(zk.getData(eq("/" + CLUSTER_NAME + "/CONTROLLER/LEADER"), anyBoolean(), any(Stat.class))) .thenReturn(json.getBytes()); InetSocketAddress address = zkManager.getControllerLeader(); assertEquals(host, address.getHostName()); assertEquals(port, address.getPort()); } @Test public void testGetClusterInfo() throws Exception { String hdfsNameNode = "namenode001"; int replicaFactor = 3; ClusterInfo clusterInfo = new ClusterInfo(hdfsNameNode, replicaFactor); when(zk.getData(eq("/" + CLUSTER_NAME), anyBoolean(), any(Stat.class))).thenReturn(clusterInfo.toJson()); ClusterInfo returnedInfo = zkManager.getClusterInfo(); assertEquals(hdfsNameNode, returnedInfo.hdfsNameNode); assertEquals(replicaFactor, returnedInfo.hdfsReplicationFactor); } @Test public void testSetClusterInfo() throws Exception { String hdfsNameNode = "namenode001"; int replicaFactor = 3; ClusterInfo clusterInfo = new ClusterInfo(hdfsNameNode, replicaFactor); when(zk.setData(eq("/" + CLUSTER_NAME), eq(clusterInfo.toJson()), anyInt())).thenReturn(new Stat()); zkManager.setClusterInfo(clusterInfo); verify(zk).setData(eq("/" + CLUSTER_NAME), eq(clusterInfo.toJson()), anyInt()); } }