Package voldemort.client.rebalance

Source Code of voldemort.client.rebalance.RebalanceSchedulerTest

/*
* Copyright 2013 LinkedIn, 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 voldemort.client.rebalance;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;

import org.junit.Test;
import org.mockito.Mockito;

import voldemort.ClusterTestUtils;
import voldemort.ServerTestUtils;
import voldemort.client.protocol.admin.AdminClient;
import voldemort.client.rebalance.task.StealerBasedRebalanceTask;
import voldemort.cluster.Cluster;

public class RebalanceSchedulerTest {

    private static List<StealerBasedRebalanceTask> sbTaskList = new ArrayList<StealerBasedRebalanceTask>();
    private static ExecutorService service = Executors.newFixedThreadPool(3);
    private static RebalanceScheduler scheduler = new RebalanceScheduler(service, 10);
    private static AdminClient adminClient;
    private final Semaphore donorPermit = new Semaphore(1);
    private final RebalanceBatchPlanProgressBar progressBar = new RebalanceBatchPlanProgressBar(0,
                                                                                                1,
                                                                                                1);

    @Test
    public void test() {

        RebalanceScheduler mockedScheduler = Mockito.spy(scheduler);
        Cluster zzCurrent = ClusterTestUtils.getZZCluster();
        adminClient = ServerTestUtils.getAdminClient(zzCurrent);

        Map<String, List<Integer>> outerMap = new HashMap<String, List<Integer>>();
        List<Integer> someList = Arrays.asList(0, 1, 2);
        outerMap.put("storeA", someList);
        int stealerId = 0;
        int donorId = 1;

        RebalanceTaskInfo partitionsInfo = new RebalanceTaskInfo(stealerId,
                                                                 donorId,
                                                                 (HashMap<String, List<Integer>>) outerMap,
                                                                 zzCurrent);

        StealerBasedRebalanceTask sbTask = new StealerBasedRebalanceTask(0,
                                                                         0,
                                                                         partitionsInfo,
                                                                         donorPermit,
                                                                         adminClient,
                                                                         progressBar,
                                                                         mockedScheduler);
        sbTaskList.add(sbTask);
        sbTaskList.add(sbTask);
        sbTaskList.add(sbTask);
        sbTaskList.add(sbTask);
        sbTaskList.add(sbTask);

        mockedScheduler.initializeLatch(sbTaskList.size());
        mockedScheduler.populateTasksByStealer(sbTaskList);

        // In the beginning both stealer and donor are idle so scheduler should
        // return the scheduled
        // task
        StealerBasedRebalanceTask scheduledTask = mockedScheduler.scheduleNextTask(false);
        org.junit.Assert.assertNotNull(sbTask);
        org.junit.Assert.assertEquals(sbTask, scheduledTask);
        mockedScheduler.removeNodesFromWorkerList(Arrays.asList(stealerId, donorId));

        // Now lets remove the donor from the worker list so that the donor is
        // idle and add the
        // stealer to the worker list. The scheduler should return null as it
        // won't be able to
        // schedule the task.
        mockedScheduler.addNodesToWorkerList(Arrays.asList(stealerId));
        mockedScheduler.removeNodesFromWorkerList(Arrays.asList(donorId));
        org.junit.Assert.assertEquals(mockedScheduler.scheduleNextTask(false), null);

        // This time stealer doesn't have any work but donor does
        mockedScheduler.addNodesToWorkerList(Arrays.asList(donorId));
        mockedScheduler.removeNodesFromWorkerList(Arrays.asList(stealerId));
        org.junit.Assert.assertEquals(mockedScheduler.scheduleNextTask(false), null);

        // And now both stealer and donor have work to do
        mockedScheduler.addNodesToWorkerList(Arrays.asList(donorId));
        mockedScheduler.addNodesToWorkerList(Arrays.asList(stealerId));
        org.junit.Assert.assertEquals(mockedScheduler.scheduleNextTask(false), null);

        // And stealer and donor are idle again
        mockedScheduler.removeNodesFromWorkerList(Arrays.asList(stealerId, donorId));
        StealerBasedRebalanceTask nextscheduledTask = mockedScheduler.scheduleNextTask(false);
        org.junit.Assert.assertNotNull(sbTask);
        org.junit.Assert.assertEquals(sbTask, nextscheduledTask);

    }
}
TOP

Related Classes of voldemort.client.rebalance.RebalanceSchedulerTest

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.