content: /**
   * 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 org.apache.aurora.scheduler;
  
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  
  import com.google.common.util.concurrent.MoreExecutors;
  import com.google.common.util.concurrent.Service;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.common.util.BackoffStrategy;
  import org.apache.aurora.scheduler.BatchWorker.Result;
  import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
  import org.apache.aurora.scheduler.testing.FakeStatsProvider;
  import org.easymock.EasyMock;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.easymock.EasyMock.expect;
  import static org.junit.Assert.assertTrue;
  
  public class BatchWorkerTest extends EasyMockTest {
    private static final String SERVICE_NAME = "TestWorker";
    private static final String BATCH_STAT = SERVICE_NAME + "_batches_processed";
    private FakeStatsProvider statsProvider;
    private BatchWorker<Boolean> batchWorker;
  
    @Before
    public void setUp() {
      StorageTestUtil storageUtil = new StorageTestUtil(this);
      storageUtil.expectOperations();
      statsProvider = new FakeStatsProvider();
      batchWorker = new BatchWorker<Boolean>(storageUtil.storage, statsProvider, 2) {
        @Override
        protected String serviceName() {
          return SERVICE_NAME;
        }
      };
    }
  
    @Test
    public void testExecute() throws Exception {
      control.replay();
  
      CompletableFuture<Boolean> result1 = batchWorker.execute(store -> true);
      CompletableFuture<Boolean> result2 = batchWorker.execute(store -> true);
      CompletableFuture<Boolean> result3 = batchWorker.execute(store -> true);
      batchWorker.startAsync().awaitRunning();
  
      assertTrue(result1.get());
      assertTrue(result2.get());
      assertTrue(result3.get());
    }
  
    @Test
    public void testExecuteThrows() throws Exception {
      control.replay();
  
      // Make sure BatchWorker service fails on unhandled error during batch processing.
      CountDownLatch shutdownLatch = new CountDownLatch(1);
      batchWorker.addListener(
          new Service.Listener() {
            @Override
            public void failed(Service.State from, Throwable failure) {
              shutdownLatch.countDown();
            }
          },
          MoreExecutors.newDirectExecutorService());
  
      batchWorker.startAsync().awaitRunning();
      batchWorker.execute(store -> {
        throw new IllegalArgumentException();
      });
  
      assertTrue(shutdownLatch.await(10L, TimeUnit.SECONDS));
    }
  
    @Test
    public void testExecuteWithReplay() throws Exception {
      BackoffStrategy backoff = createMock(BackoffStrategy.class);
      final CountDownLatch complete = new CountDownLatch(1);
  
      expect(backoff.calculateBackoffMs(EasyMock.anyLong())).andReturn(0L).anyTimes();
  
      control.replay();
  
      batchWorker.startAsync().awaitRunning();
      batchWorker.executeWithReplay(
          backoff,
          store -> statsProvider.getValue(BATCH_STAT).longValue() > 1L
              ? new Result<>(true, true)
              : new Result<>(false, false))
          .thenAccept(result -> complete.countDown());
  
      assertTrue(complete.await(10L, TimeUnit.SECONDS));
    }
  }
  