Java Code Examples for org.easymock.EasyMock#mock()

private static Message createMessage() {
    ProviderFactory factory = ServerProviderFactory.getInstance();
    Message m = new MessageImpl();
    m.put("org.apache.cxf.http.case_insensitive_queries", false);
    Exchange e = new ExchangeImpl();
    m.setExchange(e);
    e.setInMessage(m);
    Endpoint endpoint = EasyMock.mock(Endpoint.class);
    EasyMock.expect(endpoint.getEndpointInfo()).andReturn(null).anyTimes();
    EasyMock.expect(endpoint.get(Application.class.getName())).andReturn(null);
    EasyMock.expect(endpoint.size()).andReturn(0).anyTimes();
    EasyMock.expect(endpoint.isEmpty()).andReturn(true).anyTimes();
    EasyMock.expect(endpoint.get(ServerProviderFactory.class.getName())).andReturn(factory).anyTimes();
    EasyMock.replay(endpoint);
    e.put(Endpoint.class, endpoint);
    return m;
}
 

@Test
public void testForwardedRequestWithContext() {
  String scheme = "http";
  String host = "localhost";
  int port = 8443;
  String context = "/abc";
  String requestURL = String.format(Locale.ROOT, "%s://%s:%s%s", scheme, host, port, context);

  HttpServletRequest httpServletRequest = EasyMock.mock(HttpServletRequest.class);
  EasyMock.expect(httpServletRequest.getScheme()).andReturn(scheme);
  EasyMock.expect(httpServletRequest.getServerName()).andReturn(host);
  EasyMock.expect(httpServletRequest.getServerPort()).andReturn(port).anyTimes();
  EasyMock.expect(httpServletRequest.getRequestURI()).andReturn(context).anyTimes();
  EasyMock.expect(httpServletRequest.getRequestURL()).andReturn(new StringBuffer(requestURL));
  EasyMock.replay(httpServletRequest);

  String contextPath = "/mycontext";
  RequestUpdateHandler.ForwardedRequest forwardedRequest =
      new RequestUpdateHandler.ForwardedRequest(httpServletRequest, contextPath);
  Assert.assertEquals(
      String.format(Locale.ROOT, "%s://%s:%s%s", scheme, host, port, contextPath + context),
      forwardedRequest.getRequestURL().toString());
  Assert.assertEquals(contextPath + context, forwardedRequest.getRequestURI());
  Assert.assertEquals(contextPath, forwardedRequest.getContextPath());
}
 

private Message createMessage() {
    ProviderFactory factory = ServerProviderFactory.getInstance();
    Message m = new MessageImpl();
    m.put("org.apache.cxf.http.case_insensitive_queries", false);
    Exchange e = new ExchangeImpl();
    m.setExchange(e);
    e.setInMessage(m);
    Endpoint endpoint = EasyMock.mock(Endpoint.class);
    EasyMock.expect(endpoint.getEndpointInfo()).andReturn(null).anyTimes();
    EasyMock.expect(endpoint.get(Application.class.getName())).andReturn(null);
    EasyMock.expect(endpoint.size()).andReturn(0).anyTimes();
    EasyMock.expect(endpoint.isEmpty()).andReturn(true).anyTimes();
    EasyMock.expect(endpoint.get(ServerProviderFactory.class.getName())).andReturn(factory).anyTimes();
    EasyMock.replay(endpoint);
    e.put(Endpoint.class, endpoint);
    return m;
}
 

@Test
public void testReducer()
        throws IOException, InterruptedException
{
    final String expectedURI = "https://www.ukp.tu-darmstadt.de/ukp-home/";

    @SuppressWarnings("unchecked")
    final URIExtractor.URIExtractorReducer.Context context = EasyMock
            .mock(URIExtractor.URIExtractorReducer.Context.class);
    context.write(new Text(expectedURI), NullWritable.get());

    @SuppressWarnings("unchecked")
    final Iterable<NullWritable> values = EasyMock.mock(Iterable.class);

    EasyMock.replay(context, values);

    final URIExtractor.URIExtractorReducer reducer = new URIExtractor.URIExtractorReducer();

    reducer.reduce(new Text(expectedURI), values, context);

    EasyMock.verify(context, values);
}
 

private HttpServletRequest prepareRequest(HttpSession session, String userTaskId, String resource, Map<String, String []> params) {
  HttpServletRequest request = EasyMock.mock(HttpServletRequest.class);

  EasyMock.expect(request.getSession()).andReturn(session).anyTimes();
  EasyMock.expect(request.getSession(false)).andReturn(session).anyTimes();
  EasyMock.expect(request.getMethod()).andReturn(GET_METHOD).anyTimes();
  EasyMock.expect(request.getRequestURI()).andReturn(resource).anyTimes();
  EasyMock.expect(request.getParameterMap()).andReturn(params).anyTimes();
  EasyMock.expect(request.getHeader(UserTaskManager.USER_TASK_HEADER_NAME)).andReturn(userTaskId).anyTimes();
  EasyMock.expect(request.getRemoteHost()).andReturn("test-host").anyTimes();
  for (String headerName : KafkaCruiseControlServletUtils.HEADERS_TO_TRY) {
    EasyMock.expect(request.getHeader(headerName)).andReturn("localhost").anyTimes();
  }

  EasyMock.replay(request);

  return request;
}
 

private static LoadMonitor getMockLoadMonitor() {
  LoadMonitor mockLoadMonitor = EasyMock.mock(LoadMonitor.class);
  EasyMock.expect(mockLoadMonitor.taskRunnerState())
          .andReturn(LoadMonitorTaskRunner.LoadMonitorTaskRunnerState.RUNNING)
          .anyTimes();
  EasyMock.expect(mockLoadMonitor.samplingMode()).andReturn(ALL).anyTimes();
  mockLoadMonitor.pauseMetricSampling(isA(String.class), EasyMock.anyBoolean());
  expectLastCall().anyTimes();
  mockLoadMonitor.setSamplingMode(BROKER_METRICS_ONLY);
  expectLastCall().anyTimes();
  mockLoadMonitor.resumeMetricSampling(isA(String.class));
  expectLastCall().anyTimes();
  mockLoadMonitor.setSamplingMode(ALL);
  expectLastCall().anyTimes();
  return mockLoadMonitor;
}
 

/**
 * Setup the test.
 */
@BeforeClass
public static void setup() {
    final long startTime = 500L;
    MOCK_TIME = new MockTime(0, startTime, TimeUnit.NANOSECONDS.convert(startTime, TimeUnit.MILLISECONDS));
    KafkaCruiseControl mockKafkaCruiseControl = EasyMock.mock(KafkaCruiseControl.class);
    Properties props = KafkaCruiseControlUnitTestUtils.getKafkaCruiseControlProperties();
    KafkaCruiseControlConfig kafkaCruiseControlConfig = new KafkaCruiseControlConfig(props);
    EasyMock.expect(mockKafkaCruiseControl.config()).andReturn(kafkaCruiseControlConfig).anyTimes();
    EasyMock.replay(mockKafkaCruiseControl);
    Map<Integer, Long> failedBrokers = new HashMap<>();
    failedBrokers.put(1, 200L);
    failedBrokers.put(2, 400L);
    Map<String, Object> parameterConfigOverrides = new HashMap<>(4);
    parameterConfigOverrides.put(KAFKA_CRUISE_CONTROL_OBJECT_CONFIG, mockKafkaCruiseControl);
    parameterConfigOverrides.put(ANOMALY_DETECTION_TIME_MS_OBJECT_CONFIG, 200L);
    parameterConfigOverrides.put(FAILED_BROKERS_OBJECT_CONFIG, failedBrokers);
    parameterConfigOverrides.put(BROKER_FAILURES_FIXABLE_CONFIG, true);
    FAILURES = kafkaCruiseControlConfig.getConfiguredInstance(AnomalyDetectorConfig.BROKER_FAILURES_CLASS_CONFIG,
                                                              BrokerFailures.class,
                                                              parameterConfigOverrides);
}
 

private HttpServletRequest prepareRequest(HttpSession session, String userTaskId, String clientId, String resource,
    Map<String, String []> params, String method) {
  HttpServletRequest request = EasyMock.mock(HttpServletRequest.class);

  EasyMock.expect(request.getSession()).andReturn(session).anyTimes();
  EasyMock.expect(request.getSession(false)).andReturn(session).anyTimes();
  EasyMock.expect(request.getMethod()).andReturn(method).anyTimes();
  EasyMock.expect(request.getRequestURI()).andReturn(KafkaCruiseControlServletUtils.REQUEST_URI + resource).anyTimes();
  EasyMock.expect(request.getParameterMap()).andReturn(params).anyTimes();
  EasyMock.expect(request.getHeader(UserTaskManager.USER_TASK_HEADER_NAME)).andReturn(userTaskId).anyTimes();
  EasyMock.expect(request.getRemoteHost()).andReturn("test-host").anyTimes();
  for (String headerName : KafkaCruiseControlServletUtils.HEADERS_TO_TRY) {
    EasyMock.expect(request.getHeader(headerName)).andReturn(clientId).anyTimes();
  }

  for (String param : PARAMS_TO_GET) {
    String result = null;
    // Assume all parameters stored in first array entry
    if (params.get(param) != null) {
      result = params.get(param)[0];
    }
    EasyMock.expect(request.getParameter(param)).andReturn(result).anyTimes();
  }
  EasyMock.replay(request);
  return request;
}
 

@Test
public void customCompression() throws Exception
{
    final DictionaryService dictionaryService = EasyMock.mock(DictionaryService.class);

    final CompressingContentStore compressingContentStore = new CompressingContentStore();
    compressingContentStore.setNamespaceService(PREFIX_RESOLVER);
    compressingContentStore.setDictionaryService(dictionaryService);
    compressingContentStore.setCompressionType(CompressorStreamFactory.BZIP2);

    final FileContentStore fileContentStore = new FileContentStore();
    fileContentStore.setRootDirectory(backingStoreFolder.getAbsolutePath());
    fileContentStore.setProtocol("store");
    compressingContentStore.setBackingStore(fileContentStore);

    final FileContentStore temporaryContentStore = new FileContentStore();
    temporaryContentStore.setRootDirectory(temporaryStoreFolder.getAbsolutePath());
    temporaryContentStore.setProtocol("store");
    compressingContentStore.setTemporaryStore(temporaryContentStore);

    fileContentStore.afterPropertiesSet();
    temporaryContentStore.afterPropertiesSet();
    compressingContentStore.afterPropertiesSet();

    testCompressableMimetype(compressingContentStore, fileContentStore, MimetypeMap.MIMETYPE_TEXT_PLAIN, CompressorStreamFactory.BZIP2);
}
 

private void parseAttachment(String attachmentId) throws Exception {
    Attachment attachment = new AttachmentImpl(attachmentId);

    // Mock up a DataHandler for the Attachment
    DataHandler dataHandler = EasyMock.mock(DataHandler.class);
    dataHandler.setCommandMap(anyObject(CommandMap.class));
    EasyMock.expectLastCall();
    EasyMock.expect(dataHandler.getInputStream()).andReturn(null);
    EasyMock.expect(dataHandler.getContentType()).andReturn(null);
    EasyMock.replay(dataHandler);

    ((AttachmentImpl)attachment).setDataHandler(dataHandler);
    AttachmentCallbackHandler callbackHandler =
            new AttachmentCallbackHandler(Collections.singletonList(attachment));

    AttachmentRequestCallback attachmentRequestCallback = new AttachmentRequestCallback();
    attachmentRequestCallback.setAttachmentId(AttachmentUtils.getAttachmentId("cid:" + attachmentId));
    attachmentRequestCallback.setRemoveAttachments(false);

    callbackHandler.handle(new Callback[]{attachmentRequestCallback});

    List<org.apache.wss4j.common.ext.Attachment> attachments = attachmentRequestCallback.getAttachments();
    assertNotNull(attachments);
    assertEquals(1, attachments.size());
    EasyMock.verify(dataHandler);
}
 

@Test
public void testExpireSession() throws Exception {
  UUID testUserTaskId = UUID.randomUUID();

  UserTaskManager.UUIDGenerator mockUUIDGenerator = EasyMock.mock(UserTaskManager.UUIDGenerator.class);
  EasyMock.expect(mockUUIDGenerator.randomUUID()).andReturn(testUserTaskId).anyTimes();

  Time mockTime = new MockTime();
  HttpSession mockHttpSession = EasyMock.mock(HttpSession.class);
  EasyMock.expect(mockHttpSession.getLastAccessedTime()).andReturn(mockTime.milliseconds()).anyTimes();
  mockHttpSession.invalidate();

  HttpServletRequest mockHttpServletRequest = prepareRequest(mockHttpSession, null);

  OperationFuture future = new OperationFuture("future");
  UserTaskManager userTaskManager = new UserTaskManager(1000, 1, TimeUnit.HOURS.toMillis(6),
                                                        100, mockTime, mockUUIDGenerator);

  HttpServletResponse mockHttpServletResponse = EasyMock.mock(HttpServletResponse.class);
  mockHttpServletResponse.setHeader(EasyMock.anyString(), EasyMock.anyString());

  EasyMock.replay(mockUUIDGenerator, mockHttpSession, mockHttpServletResponse);
  // test-case: test if the sessions are removed on expiration
  OperationFuture future1 =
      userTaskManager.getOrCreateUserTask(mockHttpServletRequest, mockHttpServletResponse, uuid -> future, 0, true, null).get(0);
  Assert.assertEquals(future, future1);

  mockTime.sleep(1001);
  Thread.sleep(TimeUnit.SECONDS.toMillis(UserTaskManager.USER_TASK_SCANNER_PERIOD_SECONDS + 1));

  OperationFuture future2 = userTaskManager.getFuture(mockHttpServletRequest);
  Assert.assertNull(future2);

  userTaskManager.close();
}
 

@Test
public void testCompletedTasks() throws Exception {
  HttpSession mockHttpSession = EasyMock.mock(HttpSession.class);
  EasyMock.expect(mockHttpSession.getLastAccessedTime()).andReturn(100L).anyTimes();
  mockHttpSession.invalidate();

  HttpServletRequest mockHttpServletRequest = prepareRequest(mockHttpSession, null);
  UserTaskManager.UUIDGenerator mockUUIDGenerator = EasyMock.mock(UserTaskManager.UUIDGenerator.class);
  EasyMock.expect(mockUUIDGenerator.randomUUID()).andReturn(UUID.randomUUID()).anyTimes();

  OperationFuture future = new OperationFuture("future");
  UserTaskManager userTaskManager = new UserTaskManager(1000, 1, TimeUnit.HOURS.toMillis(6),
                                                        100, new MockTime(), mockUUIDGenerator);

  HttpServletResponse mockHttpServletResponse = EasyMock.mock(HttpServletResponse.class);
  Capture<String> userTaskHeader = Capture.newInstance();
  Capture<String> userTaskHeaderValue = Capture.newInstance();
  mockHttpServletResponse.setHeader(EasyMock.capture(userTaskHeader), EasyMock.capture(userTaskHeaderValue));

  EasyMock.replay(mockUUIDGenerator, mockHttpSession, mockHttpServletResponse);
  // test-case: verify if the background cleaner task removes tasks that are completed
  OperationFuture future1 =
      userTaskManager.getOrCreateUserTask(mockHttpServletRequest, mockHttpServletResponse, uuid -> future, 0, true, null).get(0);
  Assert.assertEquals(future, future1);

  future1.cancel(true);
  Thread.sleep(TimeUnit.SECONDS.toMillis(UserTaskManager.USER_TASK_SCANNER_PERIOD_SECONDS * 4));

  Assert.assertTrue(future.isDone());
  Assert.assertTrue(future.isCancelled());

  userTaskManager.close();
}
 

private static UserTasksParameters mockUserTasksParameters(HttpServletRequest answerQueryRequest) throws UnsupportedEncodingException {
  UserTasksParameters parameters = EasyMock.mock(UserTasksParameters.class);
  EasyMock.expect(parameters.userTaskIds()).andReturn(ParameterUtils.userTaskIds(answerQueryRequest)).anyTimes();
  EasyMock.expect(parameters.clientIds()).andReturn(ParameterUtils.clientIds(answerQueryRequest)).anyTimes();
  EasyMock.expect(parameters.endPoints()).andReturn(ParameterUtils.endPoints(answerQueryRequest)).anyTimes();
  EasyMock.expect(parameters.endPoint()).andReturn(ParameterUtils.endPoint(answerQueryRequest)).anyTimes();
  EasyMock.expect(parameters.types()).andReturn(ParameterUtils.types(answerQueryRequest)).anyTimes();
  EasyMock.expect(parameters.entries()).andReturn(ParameterUtils.entries(answerQueryRequest)).anyTimes();

  EasyMock.replay(parameters);
  return parameters;
}
 

@Test
public void testSamplingError() {
  KafkaCruiseControlConfig config = new KafkaCruiseControlConfig(getLoadMonitorProperties());
  Metadata metadata = new Metadata(METADATA_REFRESH_BACKOFF,
                                   METADATA_EXPIRY_MS,
                                   new LogContext(),
                                   new ClusterResourceListeners());
  MetadataClient metadataClient = new MetadataClient(config, metadata, -1L, TIME);
  MockPartitionMetricSampleAggregator mockMetricSampleAggregator =
      new MockPartitionMetricSampleAggregator(config, metadata);
  KafkaBrokerMetricSampleAggregator mockBrokerMetricSampleAggregator =
      EasyMock.mock(KafkaBrokerMetricSampleAggregator.class);
  MetricRegistry dropwizardMetricRegistry = new MetricRegistry();
  MetricSampler sampler = new MockSampler(0);
  MetricFetcherManager fetcherManager =
      new MetricFetcherManager(config, mockMetricSampleAggregator, mockBrokerMetricSampleAggregator, metadataClient,
                               METRIC_DEF, TIME, dropwizardMetricRegistry, null, sampler);
  LoadMonitorTaskRunner loadMonitorTaskRunner =
      new LoadMonitorTaskRunner(config, fetcherManager, mockMetricSampleAggregator, mockBrokerMetricSampleAggregator,
                                metadataClient, null, TIME);
  while (metadata.fetch().topics().size() < 100) {
    metadataClient.refreshMetadata();
  }
  loadMonitorTaskRunner.start(true);

  int numSamples = 0;
  long startMs = System.currentTimeMillis();
  BlockingQueue<PartitionMetricSample> sampleQueue = mockMetricSampleAggregator.metricSampleQueue();
  while (numSamples < (NUM_PARTITIONS * NUM_TOPICS) * 10 && System.currentTimeMillis() < startMs + 10000) {
    PartitionMetricSample sample = sampleQueue.poll();
    if (sample != null) {
      numSamples++;
    }
  }
  int expectedNumSamples = NUM_TOPICS * NUM_PARTITIONS;
  assertEquals("Only see " + numSamples + " samples. Expecting " + expectedNumSamples + " samples",
      expectedNumSamples, numSamples);
  fetcherManager.shutdown();
}
 

@Test
public void testInterceptorRaisesException() throws Throwable {
	EasyMock.expect(mockService.hello()).andReturn(param1);
	RequestInterceptor interceptorMock = EasyMock.mock(RequestInterceptor.class);
	interceptorMock.interceptRequest(EasyMock.anyObject(JsonNode.class));
	EasyMock.expectLastCall().andThrow(new TestThrowable(param3));
	server.setRequestInterceptor(interceptorMock);
	expectedEx.expectMessage(param3);
	expectedEx.expect(TestThrowable.class);
	EasyMock.replay(interceptorMock, mockService);
	client.hello();
	EasyMock.verify(interceptorMock, mockService);
}
 

@Test
public void mimetypeRestrictedCompression() throws Exception
{
    final DictionaryService dictionaryService = EasyMock.mock(DictionaryService.class);

    final CompressingContentStore compressingContentStore = new CompressingContentStore();
    compressingContentStore.setNamespaceService(PREFIX_RESOLVER);
    compressingContentStore.setDictionaryService(dictionaryService);
    compressingContentStore.setMimetypesToCompress(Arrays.asList(MimetypeMap.MIMETYPE_TEXT_PLAIN, MimetypeMap.MIMETYPE_XML));

    final FileContentStore fileContentStore = new FileContentStore();
    fileContentStore.setRootDirectory(backingStoreFolder.getAbsolutePath());
    fileContentStore.setProtocol("store");
    compressingContentStore.setBackingStore(fileContentStore);

    final FileContentStore temporaryContentStore = new FileContentStore();
    temporaryContentStore.setRootDirectory(temporaryStoreFolder.getAbsolutePath());
    temporaryContentStore.setProtocol("store");
    compressingContentStore.setTemporaryStore(temporaryContentStore);

    fileContentStore.afterPropertiesSet();
    temporaryContentStore.afterPropertiesSet();
    compressingContentStore.afterPropertiesSet();

    testCompressableMimetype(compressingContentStore, fileContentStore, MimetypeMap.MIMETYPE_TEXT_PLAIN, CompressorStreamFactory.GZIP);
    testCompressableMimetype(compressingContentStore, fileContentStore, MimetypeMap.MIMETYPE_XML, CompressorStreamFactory.GZIP);
    testUncompressableMimetype(compressingContentStore, fileContentStore, MimetypeMap.MIMETYPE_PDF);
}
 

@Test
public void testOnBrokerFailure() {
  final long failureTime1 = 200L;
  final long failureTime2 = 400L;
  final long startTime = 500L;
  KafkaCruiseControl mockKafkaCruiseControl = EasyMock.mock(KafkaCruiseControl.class);
  Properties props = KafkaCruiseControlUnitTestUtils.getKafkaCruiseControlProperties();
  KafkaCruiseControlConfig kafkaCruiseControlConfig = new KafkaCruiseControlConfig(props);
  EasyMock.expect(mockKafkaCruiseControl.config()).andReturn(kafkaCruiseControlConfig).atLeastOnce();
  EasyMock.replay(mockKafkaCruiseControl);
  Time mockTime = new MockTime(0, startTime, TimeUnit.NANOSECONDS.convert(startTime, TimeUnit.MILLISECONDS));
  TestingBrokerFailureAutoFixNotifier anomalyNotifier = new TestingBrokerFailureAutoFixNotifier(mockTime);
  anomalyNotifier.configure(Collections.singletonMap(SelfHealingNotifier.SELF_HEALING_BROKER_FAILURE_ENABLED_CONFIG, "true"));

  Map<Integer, Long> failedBrokers = new HashMap<>();
  failedBrokers.put(1, failureTime1);
  failedBrokers.put(2, failureTime2);
  Map<String, Object> parameterConfigOverrides = new HashMap<>(4);
  parameterConfigOverrides.put(KAFKA_CRUISE_CONTROL_OBJECT_CONFIG, mockKafkaCruiseControl);
  parameterConfigOverrides.put(FAILED_BROKERS_OBJECT_CONFIG, failedBrokers);
  parameterConfigOverrides.put(ANOMALY_DETECTION_TIME_MS_OBJECT_CONFIG, failureTime1);
  parameterConfigOverrides.put(BROKER_FAILURES_FIXABLE_CONFIG, true);

  AnomalyNotificationResult result = anomalyNotifier.onBrokerFailure(
      kafkaCruiseControlConfig.getConfiguredInstance(AnomalyDetectorConfig.BROKER_FAILURES_CLASS_CONFIG,
                                                     BrokerFailures.class,
                                                     parameterConfigOverrides));
  assertEquals(AnomalyNotificationResult.Action.CHECK, result.action());
  assertEquals(SelfHealingNotifier.DEFAULT_ALERT_THRESHOLD_MS + failureTime1 - mockTime.milliseconds(),
               result.delay());
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.BROKER_FAILURE));

  // Sleep to 1 ms before alert.
  mockTime.sleep(result.delay() - 1);
  result = anomalyNotifier.onBrokerFailure(
      kafkaCruiseControlConfig.getConfiguredInstance(AnomalyDetectorConfig.BROKER_FAILURES_CLASS_CONFIG,
                                                     BrokerFailures.class,
                                                     parameterConfigOverrides));
  assertEquals(AnomalyNotificationResult.Action.CHECK, result.action());
  assertEquals(1, result.delay());
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.BROKER_FAILURE));

  // Sleep 1 ms
  mockTime.sleep(1);
  anomalyNotifier.resetAlert(KafkaAnomalyType.BROKER_FAILURE);
  result = anomalyNotifier.onBrokerFailure(
      kafkaCruiseControlConfig.getConfiguredInstance(AnomalyDetectorConfig.BROKER_FAILURES_CLASS_CONFIG,
                                                     BrokerFailures.class,
                                                     parameterConfigOverrides));
  assertEquals(AnomalyNotificationResult.Action.CHECK, result.action());
  assertEquals(SelfHealingNotifier.DEFAULT_AUTO_FIX_THRESHOLD_MS + failureTime1 - mockTime.milliseconds(),
               result.delay());
  assertTrue(anomalyNotifier._alertCalled.get(KafkaAnomalyType.BROKER_FAILURE));

  // Sleep to 1 ms before alert.
  mockTime.sleep(result.delay() - 1);
  anomalyNotifier.resetAlert(KafkaAnomalyType.BROKER_FAILURE);
  result = anomalyNotifier.onBrokerFailure(
      kafkaCruiseControlConfig.getConfiguredInstance(AnomalyDetectorConfig.BROKER_FAILURES_CLASS_CONFIG,
                                                     BrokerFailures.class,
                                                     parameterConfigOverrides));
  assertEquals(AnomalyNotificationResult.Action.CHECK, result.action());
  assertEquals(1, result.delay());
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.BROKER_FAILURE));
  assertFalse(anomalyNotifier._autoFixTriggered.get(KafkaAnomalyType.BROKER_FAILURE));

  // Sleep 1 ms
  mockTime.sleep(1);
  anomalyNotifier.resetAlert(KafkaAnomalyType.BROKER_FAILURE);
  result = anomalyNotifier.onBrokerFailure(
      kafkaCruiseControlConfig.getConfiguredInstance(AnomalyDetectorConfig.BROKER_FAILURES_CLASS_CONFIG,
                                                     BrokerFailures.class,
                                                     parameterConfigOverrides));
  assertEquals(AnomalyNotificationResult.Action.FIX, result.action());
  assertEquals(-1L, result.delay());
  assertTrue(anomalyNotifier._alertCalled.get(KafkaAnomalyType.BROKER_FAILURE));
  assertTrue(anomalyNotifier._autoFixTriggered.get(KafkaAnomalyType.BROKER_FAILURE));
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.GOAL_VIOLATION));
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.METRIC_ANOMALY));
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.DISK_FAILURE));
  assertFalse(anomalyNotifier._alertCalled.get(KafkaAnomalyType.TOPIC_ANOMALY));
}
 

@Test
public void testSuccess(){

    PersistenceTransaction transaction = EasyMock.mock(PersistenceTransaction.class);

    // 2 times because of a redirect
    transaction.start();
    EasyMock.expectLastCall().times(2);
    transaction.end();
    EasyMock.expectLastCall().times(2);

    EasyMock.replay(transaction);

    TransactionalRequestCycleListener listener = new TransactionalRequestCycleListener(transaction);

    WicketTester tester = new WicketTester();
    tester.getApplication().getRequestCycleListeners().add(listener);

    tester.startComponentInPage(new Label("text", new Model<>()));

    EasyMock.verify(transaction);
}
 

@Test
public void testSanityCheckDryRun() throws InterruptedException, ExecutionException, TimeoutException {
  KafkaCruiseControlConfig config = EasyMock.mock(KafkaCruiseControlConfig.class);
  Time time = EasyMock.mock(Time.class);
  AnomalyDetector anomalyDetector = EasyMock.mock(AnomalyDetector.class);
  Executor executor = EasyMock.strictMock(Executor.class);
  LoadMonitor loadMonitor = EasyMock.mock(LoadMonitor.class);
  ExecutorService goalOptimizerExecutor = EasyMock.mock(ExecutorService.class);
  GoalOptimizer goalOptimizer = EasyMock.mock(GoalOptimizer.class);

  // For sanityCheckDryRun(false, true) and sanityCheckDryRun(false, false) (see #1 and #2 below).
  EasyMock.expect(executor.hasOngoingExecution()).andReturn(true).times(2);
  // For sanityCheckDryRun(false, XXX) (see #3 below)
  EasyMock.expect(executor.hasOngoingExecution()).andReturn(false).once();
  EasyMock.expect(executor.hasOngoingPartitionReassignments()).andReturn(true);
  // For sanityCheckDryRun(false, XXX) (see #4 below)
  EasyMock.expect(executor.hasOngoingExecution()).andReturn(false).once();
  EasyMock.expect(executor.hasOngoingPartitionReassignments()).andReturn(false);
  EasyMock.expect(executor.hasOngoingLeaderElection()).andReturn(true);
  // For sanityCheckDryRun(false, XXX) (see #5 below)
  EasyMock.expect(executor.hasOngoingExecution()).andReturn(false).once();
  EasyMock.expect(executor.hasOngoingPartitionReassignments()).andReturn(false);
  EasyMock.expect(executor.hasOngoingLeaderElection()).andReturn(false);

  EasyMock.replay(config, time, anomalyDetector, executor, loadMonitor, goalOptimizerExecutor, goalOptimizer);
  KafkaCruiseControl kafkaCruiseControl = new KafkaCruiseControl(config, time, anomalyDetector, executor,
                                                                 loadMonitor, goalOptimizerExecutor, goalOptimizer);

  // Expect no failure (dryrun = true) regardless of ongoing executions.
  kafkaCruiseControl.sanityCheckDryRun(true, false);
  kafkaCruiseControl.sanityCheckDryRun(true, true);
  // 1. Expect no failure (dryrun = false), if there is ongoing execution started by CC, it must be requested to stop.
  kafkaCruiseControl.sanityCheckDryRun(false, true);
  // 2. Expect failure (dryrun = false), if there is ongoing execution started by CC, not requested to stop.
  assertThrows(IllegalStateException.class, () -> kafkaCruiseControl.sanityCheckDryRun(false, false));
  // 3. Expect failure (dryrun = false), there is no execution started by CC, but ongoing replica reassignment, request to stop is irrelevant.
  assertThrows(IllegalStateException.class, () -> kafkaCruiseControl.sanityCheckDryRun(false, false));
  // 4. Expect failure (dryrun = false), there is no execution started by CC, but ongoing leader election, request to stop is irrelevant.
  assertThrows(IllegalStateException.class, () -> kafkaCruiseControl.sanityCheckDryRun(false, false));
  // 5. Expect failure (dryrun = false), there is no execution started by CC or other tools, request to stop is irrelevant.
  kafkaCruiseControl.sanityCheckDryRun(false, false);

  EasyMock.verify(config, time, anomalyDetector, executor, loadMonitor, goalOptimizerExecutor, goalOptimizer);
}
 

@Test
public void testTimeoutAndForceExecutionStop() throws InterruptedException, OngoingExecutionException {
  createTopics(0);
  // The proposal tries to move the leader. We fake the replica list to be unchanged so there is no replica
  // movement, but only leader movement.
  ExecutionProposal proposal =
      new ExecutionProposal(TP1, 0, new ReplicaPlacementInfo(1),
                            Arrays.asList(new ReplicaPlacementInfo(0), new ReplicaPlacementInfo(1)),
                            Arrays.asList(new ReplicaPlacementInfo(0), new ReplicaPlacementInfo(1)));

  KafkaCruiseControlConfig configs = new KafkaCruiseControlConfig(getExecutorProperties());
  Time time = new MockTime();
  MetadataClient mockMetadataClient = EasyMock.mock(MetadataClient.class);
  // Fake the metadata to never change so the leader movement will timeout.
  Node node0 = new Node(0, "host0", 100);
  Node node1 = new Node(1, "host1", 100);
  Node[] replicas = new Node[2];
  replicas[0] = node0;
  replicas[1] = node1;
  PartitionInfo partitionInfo = new PartitionInfo(TP1.topic(), TP1.partition(), node1, replicas, replicas);
  Cluster cluster = new Cluster("id", Arrays.asList(node0, node1), Collections.singleton(partitionInfo),
                                Collections.emptySet(), Collections.emptySet());
  MetadataClient.ClusterAndGeneration clusterAndGeneration = new MetadataClient.ClusterAndGeneration(cluster, 0);
  EasyMock.expect(mockMetadataClient.refreshMetadata()).andReturn(clusterAndGeneration).anyTimes();
  EasyMock.expect(mockMetadataClient.cluster()).andReturn(clusterAndGeneration.cluster()).anyTimes();
  LoadMonitor mockLoadMonitor = getMockLoadMonitor();
  AnomalyDetector mockAnomalyDetector = getMockAnomalyDetector(RANDOM_UUID);
  UserTaskManager.UserTaskInfo mockUserTaskInfo = getMockUserTaskInfo();
  // This tests runs two consecutive executions. First one completes w/o error, but the second one with error.
  UserTaskManager mockUserTaskManager = getMockUserTaskManager(RANDOM_UUID, mockUserTaskInfo, Arrays.asList(false, true));
  EasyMock.replay(mockMetadataClient, mockLoadMonitor, mockAnomalyDetector, mockUserTaskInfo, mockUserTaskManager);

  Collection<ExecutionProposal> proposalsToExecute = Collections.singletonList(proposal);
  Executor executor = new Executor(configs, time, new MetricRegistry(), mockMetadataClient, DEMOTION_HISTORY_RETENTION_TIME_MS,
                                   REMOVAL_HISTORY_RETENTION_TIME_MS, null, mockUserTaskManager,
                                   mockAnomalyDetector);
  executor.setExecutionMode(false);
  executor.executeProposals(proposalsToExecute,
                            Collections.emptySet(),
                            null,
                            mockLoadMonitor,
                            null,
                            null,
                            null,
                            null,
                            null,
                            null,
                            true,
                            RANDOM_UUID,
                            ExecutorTest.class::getSimpleName);
  waitUntilTrue(() -> (executor.state().state() == ExecutorState.State.LEADER_MOVEMENT_TASK_IN_PROGRESS && !executor.inExecutionTasks().isEmpty()),
                "Leader movement task did not start within the time limit",
                EXECUTION_DEADLINE_MS, EXECUTION_SHORT_CHECK_MS);

  // Sleep over ExecutorConfig#DEFAULT_LEADER_MOVEMENT_TIMEOUT_MS with some margin for inter-thread synchronization.
  time.sleep(ExecutorConfig.DEFAULT_LEADER_MOVEMENT_TIMEOUT_MS + 1L);
  // The execution should finish.
  waitUntilTrue(() -> (!executor.hasOngoingExecution() && executor.state().state() == ExecutorState.State.NO_TASK_IN_PROGRESS),
                "Proposal execution did not finish within the time limit",
                EXECUTION_DEADLINE_MS, EXECUTION_REGULAR_CHECK_MS);

  // The proposal tries to move replicas.
  proposal = new ExecutionProposal(TP1, 0, new ReplicaPlacementInfo(1),
                                   Arrays.asList(new ReplicaPlacementInfo(0), new ReplicaPlacementInfo(1)),
                                   Arrays.asList(new ReplicaPlacementInfo(1), new ReplicaPlacementInfo(0)));
  proposalsToExecute = Collections.singletonList(proposal);
  executor.executeProposals(proposalsToExecute,
                            Collections.emptySet(),
                            null,
                            mockLoadMonitor,
                            null,
                            null,
                            null,
                            null,
                            null,
                            null,
                            true,
                            RANDOM_UUID,
                            ExecutorTest.class::getSimpleName);
  waitUntilTrue(() -> (executor.state().state() == ExecutorState.State.INTER_BROKER_REPLICA_MOVEMENT_TASK_IN_PROGRESS),
                "Inter-broker replica movement task did not start within the time limit",
                EXECUTION_DEADLINE_MS, EXECUTION_SHORT_CHECK_MS);
  // Force execution to stop.
  executor.userTriggeredStopExecution(true);
  // The execution should finish.
  waitUntilTrue(() -> (!executor.hasOngoingExecution() && executor.state().state() == ExecutorState.State.NO_TASK_IN_PROGRESS),
                "Proposal execution did not finish within the time limit",
                EXECUTION_DEADLINE_MS, EXECUTION_REGULAR_CHECK_MS);
  EasyMock.verify(mockMetadataClient, mockLoadMonitor, mockAnomalyDetector, mockUserTaskInfo, mockUserTaskManager);
}