code: package edu.cmu.ml.rtw.pra.features;
  
  import java.rmi.RemoteException;
  import java.util.Arrays;
  import java.util.Map;
  import java.util.concurrent.ConcurrentHashMap;
  
  import junit.framework.TestCase;
  
  import com.google.common.collect.Maps;
  
  import edu.cmu.graphchi.walks.distributions.DiscreteDistribution;
  import edu.cmu.ml.rtw.pra.data.NodePairInstance;
  import edu.cmu.ml.rtw.pra.experiments.Outputter;
  import edu.cmu.ml.rtw.pra.graphs.GraphOnDisk;
  import com.mattg.util.FileUtil;
  import com.mattg.util.MutableConcurrentIndex;
  
  public class RandomWalkPathFinderCompanionTest extends TestCase {
    private int sourceNode = 1;
    private int targetNode = 2;
    private int intermediateNode = 3;
    private GraphOnDisk graph = new GraphOnDisk("src/test/resources/", Outputter.justLogger(), new FileUtil());
    private BasicPathTypeFactory factory = new BasicPathTypeFactory(graph);
    private PathType type1 = factory.fromString("-1-");
    private PathType type2 = factory.fromString("-2-");
    private PathType type_2 = factory.fromString("-_2-");
    private PathType type_3 = factory.fromString("-_3-");
    private PathType type12 = factory.fromString("-1-2-");
    private PathType type13 = factory.fromString("-1-3-");
    private PathType type22 = factory.fromString("-2-2-");
    private PathType type23 = factory.fromString("-2-3-");
    private MutableConcurrentIndex<PathType> pathDict = new MutableConcurrentIndex<PathType>(factory, new FileUtil());
    private int type1Index = pathDict.getIndex(type1);
    private int type2Index = pathDict.getIndex(type2);
    private int type_2Index = pathDict.getIndex(type_2);
    private int type_3Index = pathDict.getIndex(type_3);
  
    private NodePairInstance instance = new NodePairInstance(sourceNode, targetNode, true, graph);
  
    private void setDistributionsForTest(RandomWalkPathFinderCompanion companion) {
      // We'll have one source, one target, and one intermediate node in this test.
      ConcurrentHashMap<Integer, ConcurrentHashMap<Integer, DiscreteDistribution>> dists =
          new ConcurrentHashMap<Integer, ConcurrentHashMap<Integer, DiscreteDistribution>>();
      DiscreteDistribution sourceTargetDist = new DiscreteDistribution(
          new int[]{type1Index, type1Index});
      DiscreteDistribution targetSourceDist = new DiscreteDistribution(
          new int[]{type_2Index, type_2Index});
      DiscreteDistribution sourceIntermDist = new DiscreteDistribution(
          new int[]{type2Index, type2Index});
      DiscreteDistribution targetIntermDist = new DiscreteDistribution(
          new int[]{type_3Index, type_3Index});
      dists.put(sourceNode, new ConcurrentHashMap<Integer, DiscreteDistribution>());
      dists.put(targetNode, new ConcurrentHashMap<Integer, DiscreteDistribution>());
      dists.put(intermediateNode, new ConcurrentHashMap<Integer, DiscreteDistribution>());
      dists.get(targetNode).put(sourceNode, sourceTargetDist);  // will give ("-1-", 4)
      dists.get(sourceNode).put(targetNode, targetSourceDist);  // will give ("-2-", 4)
      dists.get(intermediateNode).put(sourceNode, sourceIntermDist);
      dists.get(intermediateNode).put(targetNode, targetIntermDist);  // will give ("-2-3-", 4)
      companion.setDistributions(dists);
    }
  
    public void testGetPathCounts() throws RemoteException {
      RandomWalkPathFinderCompanion companion = new RandomWalkPathFinderCompanion(
          graph, 1, 1024, new FakeVertexIdTranslate(), pathDict, factory, PathTypePolicy.PAIRED_ONLY);
      setDistributionsForTest(companion);
      Map<PathType, Integer> pathCounts = companion.getPathCounts(Arrays.asList(sourceNode),
                                                                  Arrays.asList(targetNode));
      assertEquals(3, pathCounts.size());
      assertEquals(4, pathCounts.get(type1).intValue());
      assertEquals(4, pathCounts.get(type2).intValue());
      assertEquals(4, pathCounts.get(type23).intValue());
    }
  
    public void testGetPathCountMap() throws RemoteException {
      RandomWalkPathFinderCompanion companion = new RandomWalkPathFinderCompanion(
          graph, 1, 1024, new FakeVertexIdTranslate(), pathDict, factory, PathTypePolicy.PAIRED_ONLY);
      setDistributionsForTest(companion);
      Map<NodePairInstance, Map<PathType, Integer>> pathCountMap =
          companion.getPathCountMap(Arrays.asList(instance));
      assertEquals(1, pathCountMap.size());
      assertEquals(4, pathCountMap.get(instance).get(type1).intValue());
      assertEquals(4, pathCountMap.get(instance).get(type2).intValue());
      assertEquals(4, pathCountMap.get(instance).get(type23).intValue());
    }
  
    public void testIncrementCounts() throws RemoteException {
      // First the simple case where we just have two strings and a count.
      RandomWalkPathFinderCompanion companion = new RandomWalkPathFinderCompanion(
          graph, 1, 1024, new FakeVertexIdTranslate(), pathDict, factory, PathTypePolicy.PAIRED_ONLY);
  
      Map<PathType, Integer> pathCounts = Maps.newHashMap();
      companion.incrementCounts(pathCounts, type1, type_2, 2);
      assertEquals(2, pathCounts.get(type12).intValue());
      companion.incrementCounts(pathCounts, type1, type_2, 1);
      assertEquals(3, pathCounts.get(type12).intValue());
  
      // Now assume we have a distribution and a path.  The distribution here is a distribution
      // of _path types_ that connect two nodes.
      DiscreteDistribution dist = new DiscreteDistribution(
          new int[]{type_2Index, type_2Index, type_3Index});
      pathCounts = Maps.newHashMap();
      companion.incrementCounts(pathCounts, type1, dist);
      assertEquals(4, pathCounts.get(type12).intValue());
      assertEquals(1, pathCounts.get(type13).intValue());
  
      pathCounts = Maps.newHashMap();
      companion.setPolicy(PathTypePolicy.EVERYTHING);
      companion.incrementCounts(pathCounts, type1, dist);
      assertEquals(8, pathCounts.get(type12).intValue());
      assertEquals(1, pathCounts.get(type13).intValue());
  
      dist = new DiscreteDistribution( new int[]{type1Index, type1Index, type2Index});
      pathCounts = Maps.newHashMap();
      companion.incrementCounts(pathCounts, dist, type_2);
      assertEquals(8, pathCounts.get(type12).intValue());
      assertEquals(1, pathCounts.get(type22).intValue());
  
      pathCounts = Maps.newHashMap();
      companion.setPolicy(PathTypePolicy.PAIRED_ONLY);
      companion.incrementCounts(pathCounts, dist, type_2);
      assertEquals(4, pathCounts.get(type12).intValue());
      assertEquals(1, pathCounts.get(type22).intValue());
  
      // Now we have two distributions that we're combining, one from the source, one from the
      // target.
      DiscreteDistribution sourceDist = new DiscreteDistribution(
          new int[]{type1Index, type1Index, type2Index});
      DiscreteDistribution targetDist = new DiscreteDistribution(
          new int[]{type_2Index, type_2Index, type_3Index});
      pathCounts = Maps.newHashMap();
      companion.incrementCounts(pathCounts, sourceDist, targetDist);
      assertEquals(4, pathCounts.get(type12).intValue());
      assertEquals(2, pathCounts.get(type13).intValue());
      assertEquals(2, pathCounts.get(type22).intValue());
      assertEquals(1, pathCounts.get(type23).intValue());
    }
  }
  