@Before
public void setUp() {
    stackName = UUID.randomUUID().toString();

    GrpcConfiguration grpcClientConfiguration = mock(GrpcConfiguration.class);
    when(grpcClientConfiguration.getRequestTimeoutMs()).thenReturn(1000L);

    TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class);
    when(titusFederationConfiguration.getStack()).thenReturn(stackName);
    when(titusFederationConfiguration.getCells()).thenReturn("one=1");
    when(titusFederationConfiguration.getRoutingRules()).thenReturn("one=(app1.*|app2.*);two=(app3.*)");

    CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration);
    ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration);
    List<Cell> cells = cellInfoResolver.resolve();
    cellToServiceMap = ImmutableMap.of(cells.get(0), cell);

    CellConnector connector = mock(CellConnector.class);
    when(connector.getChannels()).thenReturn(cellToServiceMap.entrySet().stream()
            .collect(Collectors.toMap(Map.Entry::getKey, cellPairEntry -> cellPairEntry.getValue().getChannel()))
    );
    when(connector.getChannelForCell(any(Cell.class))).thenAnswer(invocation ->
            Optional.ofNullable(cellToServiceMap.get(invocation.<Cell>getArgument(0)))
                    .map(GrpcServerRule::getChannel)
    );

    final AggregatingCellClient aggregatingCellClient = new AggregatingCellClient(connector);
    final AnonymousCallMetadataResolver anonymousCallMetadataResolver = new AnonymousCallMetadataResolver();
    service = new AggregatingJobServiceGateway(
            grpcClientConfiguration,
            titusFederationConfiguration,
            connector,
            cellRouter,
            aggregatingCellClient,
            new AggregatingJobManagementServiceHelper(aggregatingCellClient, grpcClientConfiguration)
    );

    clock = Clocks.test();
    dataGenerator = new ServiceDataGenerator(clock, TASKS_IN_GENERATED_JOBS);
}
 

@Before
public void setUp() {
    stackName = UUID.randomUUID().toString();

    GrpcConfiguration grpcConfiguration = mock(GrpcConfiguration.class);
    when(grpcConfiguration.getRequestTimeoutMs()).thenReturn(GRPC_REQUEST_TIMEOUT_MS);

    TitusFederationConfiguration titusFederationConfiguration = mock(TitusFederationConfiguration.class);
    when(titusFederationConfiguration.getStack()).thenReturn(stackName);
    when(titusFederationConfiguration.getCells()).thenReturn("one=1;two=2");
    when(titusFederationConfiguration.getRoutingRules()).thenReturn("one=(app1.*|app2.*);two=(app3.*)");

    CellInfoResolver cellInfoResolver = new DefaultCellInfoResolver(titusFederationConfiguration);
    ApplicationCellRouter cellRouter = new ApplicationCellRouter(cellInfoResolver, titusFederationConfiguration);
    cells = cellInfoResolver.resolve();
    cellToServiceMap = ImmutableMap.of(
            cells.get(0), cellOne,
            cells.get(1), cellTwo
    );

    CellConnector connector = mock(CellConnector.class);
    when(connector.getChannels()).thenReturn(cellToServiceMap.entrySet().stream()
            .collect(Collectors.toMap(Map.Entry::getKey, cellPairEntry -> cellPairEntry.getValue().getChannel()))
    );
    when(connector.getChannelForCell(any(Cell.class))).thenAnswer(invocation ->
            Optional.ofNullable(cellToServiceMap.get(invocation.<Cell>getArgument(0)))
                    .map(GrpcServerRule::getChannel)
    );

    final AggregatingCellClient aggregatingCellClient = new AggregatingCellClient(connector);
    service = new AggregatingJobServiceGateway(
            grpcConfiguration,
            titusFederationConfiguration,
            connector,
            cellRouter,
            aggregatingCellClient,
            new AggregatingJobManagementServiceHelper(aggregatingCellClient, grpcConfiguration)
    );

    clock = Clocks.test();
    dataGenerator = new ServiceDataGenerator(clock, TASKS_IN_GENERATED_JOBS);
}