package org.apache.helix;
   
   /*
    * 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.
   */
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  import java.util.Map;
  import java.util.Random;
  import java.util.Set;
  import java.util.TreeMap;
  
  import org.apache.helix.ZNRecord;
  import org.apache.helix.controller.ExternalViewGenerator;
  import org.apache.helix.model.Message;
  import org.apache.helix.model.CurrentState.CurrentStateProperty;
  import org.testng.AssertJUnit;
  import org.testng.annotations.Test;
  
  
  public class TestZKRoutingInfoProvider
  {
    public Map<String, List<ZNRecord>> createCurrentStates(String[] dbNames,
        String[] nodeNames, int[] partitions, int[] replicas)
    {
      Map<String, List<ZNRecord>> currentStates = new TreeMap<String, List<ZNRecord>>();
      Map<String, Map<String, ZNRecord>> currentStates2 = new TreeMap<String, Map<String, ZNRecord>>();
  
      Map<String, String> stateMaster = new TreeMap<String, String>();
      stateMaster.put(CurrentStateProperty.CURRENT_STATE.toString(), "MASTER");
  
      Map<String, String> stateSlave = new TreeMap<String, String>();
      stateSlave.put(CurrentStateProperty.CURRENT_STATE.toString(), "SLAVE");
  
      for (int i = 0; i < nodeNames.length; i++)
      {
        currentStates.put(nodeNames[i], new ArrayList<ZNRecord>());
        currentStates2.put(nodeNames[i], new TreeMap<String, ZNRecord>());
        for (int j = 0; j < dbNames.length; j++)
        {
          ZNRecord dbPartitionState = new ZNRecord(dbNames[j]);
          currentStates2.get(nodeNames[i]).put(dbNames[j], dbPartitionState);
        }
      }
  
      Random rand = new Random(1234);
      for (int j = 0; j < dbNames.length; j++)
      {
        int partition = partitions[j];
        ArrayList<Integer> randomArray = new ArrayList<Integer>();
        for (int i = 0; i < partition; i++)
        {
          randomArray.add(i);
        }
        Collections.shuffle(randomArray, rand);
  
        for (int i = 0; i < partition; i++)
        {
          stateMaster.put(Message.Attributes.RESOURCE_NAME.toString(),
              dbNames[j]);
          stateSlave.put(Message.Attributes.RESOURCE_NAME.toString(),
              dbNames[j]);
          int nodes = nodeNames.length;
          int master = randomArray.get(i) % nodes;
          String partitionName = dbNames[j] + ".partition-" + i;
          Map<String, Map<String, String>> map = (currentStates2
              .get(nodeNames[master]).get(dbNames[j]).getMapFields());
          assert (map != null);
          map.put(partitionName, stateMaster);
  
          for (int k = 1; k <= replicas[j]; k++)
          {
            int slave = (master + k) % nodes;
            Map<String, Map<String, String>> map2 = currentStates2
                .get(nodeNames[slave]).get(dbNames[j]).getMapFields();
  
            map2.put(partitionName, stateSlave);
          }
        }
      }
      for (String nodeName : currentStates2.keySet())
      {
       Map<String, ZNRecord> recMap = currentStates2.get(nodeName);
       List<ZNRecord> list = new ArrayList<ZNRecord>();
       for (ZNRecord rec : recMap.values())
       {
         list.add(rec);
       }
       currentStates.put(nodeName, list);
     }
     return currentStates;
   }
 
   private void verify(Map<String, List<ZNRecord>> currentStates,
       Map<String, Map<String, Set<String>>> routingMap)
   {
     int counter1 = 0;
     int counter2 = 0;
     for (String nodeName : currentStates.keySet())
     {
       List<ZNRecord> dbStateList = currentStates.get(nodeName);
       for (ZNRecord dbState : dbStateList)
       {
         Map<String, Map<String, String>> dbStateMap = dbState.getMapFields();
         for (String partitionName : dbStateMap.keySet())
         {
           Map<String, String> stateMap = dbStateMap
               .get(partitionName);
           String state = stateMap
               .get(CurrentStateProperty.CURRENT_STATE.toString());
           AssertJUnit.assertTrue(routingMap.get(partitionName).get(state)
               .contains(nodeName));
           counter1++;
         }
       }
     }
 
     for (String partitionName : routingMap.keySet())
     {
       Map<String, Set<String>> partitionState = routingMap.get(partitionName);
       for (String state : partitionState.keySet())
       {
         counter2 += partitionState.get(state).size();
       }
     }
     AssertJUnit.assertTrue(counter2 == counter1);
   }
 
   // public static void main(String[] args)
   @Test ()
   public void testInvocation() throws Exception
   {
     String[] dbNames = new String[3];
     for (int i = 0; i < dbNames.length; i++)
     {
       dbNames[i] = "DB_" + i;
     }
     String[] nodeNames = new String[6];
     for (int i = 0; i < nodeNames.length; i++)
     {
       nodeNames[i] = "LOCALHOST_100" + i;
     }
 
     int[] partitions = new int[dbNames.length];
     for (int i = 0; i < partitions.length; i++)
     {
       partitions[i] = (i + 1) * 10;
     }
 
     int[] replicas = new int[dbNames.length];
     for (int i = 0; i < replicas.length; i++)
     {
       replicas[i] = 3;
     }
     Map<String, List<ZNRecord>> currentStates = createCurrentStates(dbNames,
         nodeNames, partitions, replicas);
     ExternalViewGenerator provider = new ExternalViewGenerator();
 
     List<ZNRecord> mockIdealStates = new ArrayList<ZNRecord>();
     for (String dbName : dbNames)
     {
       ZNRecord rec = new ZNRecord(dbName);
       mockIdealStates.add(rec);
     }
     List<ZNRecord> externalView = provider.computeExternalView(currentStates,
         mockIdealStates);
 
     Map<String, Map<String, Set<String>>> routingMap = provider
         .getRouterMapFromExternalView(externalView);
 
     verify(currentStates, routingMap);
 
     /* write current state and external view to ZK */
     /*
      * String clusterName = "test-cluster44"; ZkClient zkClient = new
      * ZkClient("localhost:2181"); zkClient.setZkSerializer(new
      * ZNRecordSerializer());
      *
      * for(String nodeName : currentStates.keySet()) {
      * if(zkClient.exists(CMUtil.getCurrentStatePath(clusterName, nodeName))) {
      * zkClient.deleteRecursive(CMUtil.getCurrentStatePath(clusterName,
      * nodeName)); } ZKUtil.createChildren(zkClient,CMUtil.getCurrentStatePath
      * (clusterName, nodeName), currentStates.get(nodeName)); }
      *
      * //List<ZNRecord> externalView =
      * ZKRoutingInfoProvider.computeExternalView(currentStates); String
      * routingTablePath = CMUtil.getExternalViewPath(clusterName);
      * if(zkClient.exists(routingTablePath)) {
      * zkClient.deleteRecursive(routingTablePath); }
      *
      * ZKUtil.createChildren(zkClient, CMUtil.getExternalViewPath(clusterName),
      * externalView);
      */
   }
 }