private final Wrapper newSplitWrapper(double cost, int minWidth, int maxWidth, List<EndpointAffinity> epAffs,
ExecutionNodeMap execNodeMap) {
PhysicalOperator root = Mockito.mock(PhysicalOperator.class);
Fragment fragment = Mockito.mock(Fragment.class);
when(fragment.getRoot()).thenReturn(root);
final Wrapper fragmentWrapper = new Wrapper(fragment, 1);
final Stats stats = fragmentWrapper.getStats();
stats.addCost(cost);
stats.addMinWidth(minWidth);
stats.addMaxWidth(maxWidth);
DatasetSplit dataSplit = new DatasetSplit();
Iterable<Affinity> affinities = Iterables.transform(
epAffs,
new Function<EndpointAffinity, Affinity>() {
@Override
public Affinity apply(@Nullable EndpointAffinity endpointAffinity) {
return new Affinity()
.setFactor(endpointAffinity.getAffinity())
.setHost(endpointAffinity.getEndpoint()
.getAddress());
}
}
);
dataSplit.setAffinitiesList(Lists.newArrayList(affinities));
SplitWork splitWork = new SplitWork(dataSplit, execNodeMap, DistributionAffinity.HARD);
GroupScan groupScan = Mockito.mock(GroupScan.class);
when(groupScan.getDistributionAffinity()).thenReturn(DistributionAffinity.HARD);
stats.addSplits(groupScan, ImmutableList.<CompleteWork>of(splitWork));
return fragmentWrapper;
}