TerrapinUtil { public static void loadFileSetData(ZooKeeperManager zkManager, FileSetInfo fsInfo, Options options) throws Exception { InetSocketAddress controllerSockAddress = zkManager.getControllerLeader(); LOG.info("Connecting to controller at " + controllerSockAddress.getHostName() + ":" + controllerSockAddress.getPort()); LOG.info("Load timeout " + Constants.LOAD_TIMEOUT_SECONDS + " seconds."); Service<ThriftClientRequest, byte[]> service = ClientBuilder.safeBuild(ClientBuilder.get() .hosts(controllerSockAddress) .codec(new ThriftClientFramedCodecFactory(Option.<ClientId>empty())) .retries(1) .connectTimeout(Duration.fromMilliseconds(1000)) .requestTimeout(Duration.fromSeconds(Constants.LOAD_TIMEOUT_SECONDS)) .hostConnectionLimit(100) .failFast(false)); TerrapinController.ServiceIface iface = new TerrapinController.ServiceToClient( service, new TBinaryProtocol.Factory()); TerrapinLoadRequest request = new TerrapinLoadRequest(); request.setHdfsDirectory(fsInfo.servingInfo.hdfsPath); request.setOptions(options); request.setFileSet(fsInfo.fileSetName); request.setExpectedNumPartitions(fsInfo.servingInfo.numPartitions); LOG.info("Loading file set " + fsInfo.fileSetName + " at " + fsInfo.servingInfo.hdfsPath); long startTimeSeconds = System.currentTimeMillis() / 1000; int numTriesLeft = 5; boolean done = false; Exception e = null; while (numTriesLeft > 0) { try { iface.loadFileSet(request).get(); done = true; break; } catch (Throwable t) { LOG.error("Swap failed with exception.", t); e = new Exception(t); numTriesLeft--; } LOG.info("Retrying in 10 seconds."); try { Thread.sleep(10000); } catch (InterruptedException ie) { LOG.error("Interrupted."); break; } } if (done) { LOG.info("Load successful. Swap took " + ((System.currentTimeMillis() / 1000) - startTimeSeconds) + " seconds."); } else { LOG.error("Load failed !!!."); throw new Exception(e); } }  static String getHelixInstanceFromHDFSHost(String hdfsHostName); static PropertiesConfiguration readPropertiesExitOnFailure(String configFile); static Integer extractPartitionName(String fileName, PartitionerType partitioner); static String formatPartitionName(int partitionNumber); static String getPartitionName(ByteBuffer key,
                                        PartitionerType partitionerType,
                                        int numPartitions); static String hdfsDirToHelixResource(String hdfsDir); static String helixResourceToHdfsDir(String helixResource); static ZooKeeperClient getZooKeeperClient(String zkQuorum, int sessionTimeoutSeconds); static int getBucketSize(int numPartitions, boolean enableZkCompression); static String extractFileSetFromPath(String resource); static Pair<String, Integer> getBucketizedResourceAndPartitionNum(String helixPartition); static Pair<String, Integer> getNonBucketizedResourceAndPartitionNum(
      String helixPartition); static Pair<String, Integer> getResourceAndPartitionNum(String helixPartition); static String getViewPartitionName(String resource, int partition); static int getViewPartitionNumber(String viewPartitionName); static void compressIdealState(IdealState is); static String getZKQuorumFromConf(PropertiesConfiguration configuration); static List<HdfsFileStatus> getHdfsFileList(DFSClient hdfsClient,
                                                     String hdfsDir); static void loadFileSetData(ZooKeeperManager zkManager, FileSetInfo fsInfo, Options options); static List<Pair<Path, Long>> getS3FileList(AWSCredentials credentials,
      String s3Bucket, String s3KeyPrefix); static void setupConfiguration(Configuration conf,
                                        long dfsBlockSize,
                                        int dfsReplication);  }
@Test @PrepareForTest(TerrapinUtil.class) public void testLoadFileSetData() throws Exception { TerrapinController.ServiceToClient serviceImpl = mock(TerrapinController.ServiceToClient.class); ZooKeeperManager zkManager = mock(ZooKeeperManager.class); FileSetInfo fileSetInfo = new FileSetInfo("fileset", "dir", 100, new ArrayList<FileSetInfo.ServingInfo>(), new Options()); InetSocketAddress address = new InetSocketAddress("namenode001", 9090); Options requestOptions = new Options(); when(zkManager.getControllerLeader()).thenReturn(address); whenNew(TerrapinController.ServiceToClient.class).withAnyArguments().thenReturn(serviceImpl); when(serviceImpl.loadFileSet(any(TerrapinLoadRequest.class))).thenReturn(Future.Void()); ArgumentCaptor<TerrapinLoadRequest> requestCaptor = ArgumentCaptor.forClass(TerrapinLoadRequest.class); TerrapinUtil.loadFileSetData(zkManager, fileSetInfo, requestOptions); verify(serviceImpl).loadFileSet(requestCaptor.capture()); TerrapinLoadRequest request = requestCaptor.getValue(); assertEquals(fileSetInfo.servingInfo.hdfsPath, request.getHdfsDirectory()); assertEquals(requestOptions, request.getOptions()); assertEquals(fileSetInfo.fileSetName, request.getFileSet()); assertEquals(fileSetInfo.servingInfo.numPartitions, request.getExpectedNumPartitions()); }