use of com.linkedin.darkcluster.api.DarkRequestHeaderGenerator in project rest.li by linkedin.
the class TestDarkClusterManager method testWithDarkHeaders.
@Test
public void testWithDarkHeaders() {
MockClusterInfoProvider clusterInfoProvider = new MockClusterInfoProvider();
Facilities facilities = new MockFacilities(clusterInfoProvider);
// This configuration will choose the RelativeTrafficMultiplierDarkClusterStrategy
DarkClusterConfig darkClusterConfig = createRelativeTrafficMultiplierConfig(1.0f);
clusterInfoProvider.addDarkClusterConfig(SOURCE_CLUSTER_NAME, DARK_CLUSTER_NAME, darkClusterConfig);
DarkClusterStrategyFactory mockStrategyFactory = Mockito.mock(DarkClusterStrategyFactory.class);
DarkClusterStrategy mockDarkStrategy = Mockito.mock(DarkClusterStrategy.class);
DarkRequestHeaderGenerator darkRequestHeaderGenerator = Mockito.mock(DarkRequestHeaderGenerator.class);
Mockito.when(mockStrategyFactory.get(DARK_CLUSTER_NAME)).thenReturn(mockDarkStrategy);
Mockito.when(darkRequestHeaderGenerator.get(DARK_CLUSTER_NAME)).thenReturn(Optional.of(new DarkRequestHeaderGenerator.HeaderNameValuePair("header", "value")));
RestRequest restRequest = new RestRequestBuilder(URI.create("/abc")).setMethod(METHOD_SAFE).build();
RestRequest darkRequest = new RestRequestBuilder(URI.create("d2://" + DARK_CLUSTER_NAME + "/abc")).setMethod(METHOD_SAFE).setHeader("header", "value").build();
RequestContext requestContext = new RequestContext();
Mockito.when(mockDarkStrategy.handleRequest(restRequest, darkRequest, new RequestContext(requestContext))).thenReturn(true);
DarkClusterManager darkClusterManager = new DarkClusterManagerImpl(SOURCE_CLUSTER_NAME, facilities, mockStrategyFactory, null, null, new DoNothingNotifier(), null, Collections.singletonList(darkRequestHeaderGenerator));
boolean status = darkClusterManager.handleDarkRequest(restRequest, requestContext);
Assert.assertTrue(status);
}
Aggregations