/*
 * Copyright 2018 LinkedIn Corp.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */

package azkaban.utils;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.catchThrowable;

import java.time.Duration;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.RejectedExecutionException;
import org.junit.Test;

@SuppressWarnings("FutureReturnValueIgnored")
public class ExecutorServiceUtilsTest {

  private final ExecutorServiceUtils executorServiceUtils = new ExecutorServiceUtils();

  private boolean isTestThreadInterrupted = false;

  @Test
  public void gracefulShutdown() throws InterruptedException {
    // given
    final ExecutorService service = Executors.newSingleThreadExecutor();

    // when
    this.executorServiceUtils.gracefulShutdown(service, Duration.ofMillis(1));

    // then
    assertThat(service.isShutdown()).isTrue();
  }

  @Test
  public void force_shutdown_after_timeout() throws InterruptedException {
    // given
    final ExecutorService service = Executors.newSingleThreadExecutor();

    // when
    service.submit(this::sleep);
    final long beginShutdownTime = System.currentTimeMillis();
    this.executorServiceUtils.gracefulShutdown(service, Duration.ofMillis(1));
    final long endShutdownTime = System.currentTimeMillis();

    // then
    assertThat(service.isShutdown()).isTrue();
    assertThat(this.isTestThreadInterrupted).isTrue();
    final long shutdownDuration = endShutdownTime - beginShutdownTime;

    // Test that the wait time parameter is effective.
    // Give some buffer to account for overhead to reduce false positives.
    assertThat(shutdownDuration).isLessThan(1000);
  }

  @Test
  public void can_not_submit_tasks_after_shutdown() throws
      InterruptedException {
    // given
    final ExecutorService service = Executors.newSingleThreadExecutor();
    service.submit(this::sleep);

    // when
    this.executorServiceUtils.gracefulShutdown(service, Duration.ofMillis(1));
    final Throwable thrown = catchThrowable(() -> service.submit(this::sleep));

    // then
    assertThat(thrown).isInstanceOf(RejectedExecutionException.class);
  }

  private void sleep() {
    try {
      Thread.sleep(5000);
    } catch (final InterruptedException ex) {
      this.isTestThreadInterrupted = true;
    }
  }
}

@@ -731,16 +731,138 @@
 rowable;
+%0Aimport static org.mockito.Mockito.mock;%0Aimport static org.mockito.Mockito.verify;%0Aimport static org.mockito.Mockito.when;
 %0A%0Aimport
@@ -1018,24 +1018,62 @@
 nException;%0A
+import java.util.concurrent.TimeUnit;%0A
 import org.j
@@ -1264,60 +1264,8 @@
 );%0A%0A
-  private boolean isTestThreadInterrupted = false;%0A%0A
   @T
@@ -1702,32 +1702,37 @@
 e service = 
+mock(
 Executor
 s.newSingleT
@@ -1723,62 +1723,41 @@
 utor
-s.newSingleThreadExecutor(
+Service.class
 );%0A
-%0A 
    
-//
  when
-%0A    
+(
 service.
 subm
@@ -1756,91 +1756,83 @@
 ice.
-submit(this::sleep);%0A    final long beginShutdownTime = System.currentTimeMillis();
+awaitTermination(1, TimeUnit.MILLISECONDS)).thenReturn(false);%0A%0A    // when
 %0A   
@@ -1910,69 +1910,8 @@
 1));
-%0A    final long endShutdownTime = System.currentTimeMillis();
 %0A%0A  
@@ -1924,34 +1924,30 @@
 hen%0A    
-assertThat
+verify
 (service
 .isShutd
@@ -1930,36 +1930,35 @@
   verify(service
-.isS
+).s
 hutdown()).isTru
@@ -1954,325 +1954,42 @@
 wn()
-).isTrue();%0A    assertThat(this.isTestThreadInterrupted).isTrue();%0A    final long shutdownDuration = endShutdownTime - beginShutdownTime;%0A%0A    // Test that the wait time parameter is effective.%0A    // Give some buffer to account for overhead to reduce false positives.%0A    assertThat(shutdownDuration).isLessThan(1000
+;%0A    verify(service).shutdownNow(
 );%0A 
@@ -2579,51 +2579,8 @@
 ) %7B%0A
-      this.isTestThreadInterrupted = true;%0A
     

Fix flaky test force_shutdown_after_timeout (#1844)