1/*2 * Licensed to the Apache Software Foundation (ASF) under one3 * or more contributor license agreements. See the NOTICE file4 * distributed with this work for additional information5 * regarding copyright ownership. The ASF licenses this file6 * to you under the Apache License, Version 2.0 (the7 * "License"); you may not use this file except in compliance8 * with the License. You may obtain a copy of the License at9 *10 * http://www.apache.org/licenses/LICENSE-2.011 *12 * Unless required by applicable law or agreed to in writing, software13 * distributed under the License is distributed on an "AS IS" BASIS,14 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15 * See the License for the specific language governing permissions and16 * limitations under the License.17 */18package org.apache.giraph.partition;
1920import org.apache.giraph.edge.ArrayListEdges;
21import org.apache.giraph.edge.EdgeFactory;
22import org.apache.giraph.edge.OutEdges;
23import org.apache.giraph.graph.DefaultVertex;
24import org.apache.giraph.graph.GiraphTransferRegulator;
25import org.apache.giraph.graph.Vertex;
26import org.apache.giraph.job.GiraphJob;
27import org.apache.giraph.utils.NoOpComputation;
28import org.apache.hadoop.io.DoubleWritable;
29import org.apache.hadoop.io.FloatWritable;
30import org.apache.hadoop.io.IntWritable;
31import org.apache.hadoop.io.LongWritable;
32import org.junit.Before;
33import org.junit.Test;
3435import java.io.IOException;
3637importstatic org.junit.Assert.assertFalse;
38importstatic org.junit.Assert.assertTrue;
39importstatic org.mockito.Mockito.mock;
40importstatic org.mockito.Mockito.when;
4142/**43 * Test the GiraphTransferRegulator.44 */45publicclassTestGiraphTransferRegulator {
46/** Job filled in by setup() */47private GiraphJob job;
48/** Instantiated vertex filled in from setup() */49private Vertex<IntWritable, FloatWritable, DoubleWritable>
50 vertex = new DefaultVertex<IntWritable, FloatWritable, DoubleWritable>();
5152/**53 * Dummy vertex.54 */55publicstaticclassTestComputationextends NoOpComputation<IntWritable,
56 FloatWritable, DoubleWritable, LongWritable> { }
5758 @Before
59publicvoid setUp() {
60try {
61 job = new GiraphJob("TestGiraphTransferRegulator");
62 } catch (IOException e) {
63thrownew RuntimeException("setUp: Failed", e);
64 }
65 job.getConfiguration().setComputationClass(TestComputation.class);
66 }
6768 @Test
69publicvoid testGiraphTransferRegulator() {
70 job.getConfiguration()
71 .setInt(GiraphTransferRegulator.MAX_VERTICES_PER_TRANSFER, 1);
72 job.getConfiguration()
73 .setInt(GiraphTransferRegulator.MAX_EDGES_PER_TRANSFER, 3);
74 OutEdges<IntWritable, DoubleWritable> edges =
75new ArrayListEdges<IntWritable, DoubleWritable>();
76 edges.initialize(3);
77 edges.add(EdgeFactory.create(new IntWritable(2), new DoubleWritable(22)));
78 edges.add(EdgeFactory.create(new IntWritable(3), new DoubleWritable(33)));
79 edges.add(EdgeFactory.create(new IntWritable(4), new DoubleWritable(44)));
80 vertex.initialize(new IntWritable(1), new FloatWritable(1), edges);
81 GiraphTransferRegulator gtr =
82new GiraphTransferRegulator(job.getConfiguration());
83 PartitionOwner owner = mock(PartitionOwner.class);
84 when(owner.getPartitionId()).thenReturn(57);
85 assertFalse(gtr.transferThisPartition(owner));
86 gtr.incrementCounters(owner, vertex);
87 assertTrue(gtr.transferThisPartition(owner));
88 }
8990 }