1 // Copyright 2021 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 package com.google.android.downloader;
16 
17 import static com.google.common.base.Throwables.getStackTraceAsString;
18 import static com.google.common.truth.Truth.assertThat;
19 import static java.util.concurrent.TimeUnit.MILLISECONDS;
20 import static java.util.stream.Collectors.joining;
21 import static org.junit.Assert.fail;
22 
23 import java.time.Duration;
24 import java.util.ArrayList;
25 import java.util.List;
26 import java.util.concurrent.ExecutorService;
27 import java.util.concurrent.Executors;
28 import java.util.concurrent.ScheduledExecutorService;
29 import java.util.concurrent.ThreadFactory;
30 import org.junit.rules.ExternalResource;
31 
32 /**
33  * A {@link org.junit.rules.TestRule} that manages and provides instances of {@link
34  * java.util.concurrent.Executor} and its various more specific interfaces. Takes care of shutting
35  * down any started threads and executors during execution, and also collects uncaught exceptions,
36  * failing the test and reporting the uncaught exception if any are found during execution.
37  */
38 public class TestExecutorRule extends ExternalResource {
39   private final Duration timeout;
40   private final List<Throwable> uncaughtExceptions = new ArrayList<>();
41   private final List<ExecutorService> executorServices = new ArrayList<>();
42   private final ThreadFactory threadFactory =
43       runnable -> {
44         Thread thread = Executors.defaultThreadFactory().newThread(runnable);
45         // Insert an uncaught exception handler so that that errors happening on a background
46         // thread can be collected and cause test failures.
47         thread.setUncaughtExceptionHandler((t, e) -> uncaughtExceptions.add(e));
48         return thread;
49       };
50 
51   /**
52    * Constructs a new instance of this rule with the provided {@code timeout}. The timeout will be
53    * used when calling {@link ExecutorService#awaitTermination} on any {@link ExecutorService}
54    * instances created by this rule.
55    */
TestExecutorRule(Duration timeout)56   public TestExecutorRule(Duration timeout) {
57     this.timeout = timeout;
58   }
59 
60   /**
61    * Creates a new single-threaded {@link ExecutorService} for use in tests. The executor will
62    * collect any uncaught exceptions encountered during test execution, and will fail the test with
63    * a detailed report of exceptions, if any are encountered. The executor will also be shut down
64    * and will await termination. Failure to shutdown in time (e.g. due to a blocked thread) will
65    * result in a test failure as well.
66    */
newSingleThreadExecutor()67   public ExecutorService newSingleThreadExecutor() {
68     ExecutorService executorService = Executors.newSingleThreadExecutor(threadFactory);
69     executorServices.add(executorService);
70     return executorService;
71   }
72 
73   /**
74    * Creates a new single-threaded {@link ScheduledExecutorService} for use in tests. The executor
75    * will collect any uncaught exceptions encountered during test execution, and will fail the test
76    * with a detailed report of exceptions, if any are encountered. The executor will also be shut
77    * down and will await termination. Failure to shutdown in time (e.g. due to a blocked thread)
78    * will result in a test failure as well.
79    */
newSingleThreadScheduledExecutor()80   public ScheduledExecutorService newSingleThreadScheduledExecutor() {
81     ScheduledExecutorService executorService =
82         Executors.newSingleThreadScheduledExecutor(threadFactory);
83     executorServices.add(executorService);
84     return executorService;
85   }
86 
87   @Override
after()88   protected void after() {
89     try {
90       for (ExecutorService executorService : executorServices) {
91         try {
92           executorService.shutdown();
93           assertThat(executorService.awaitTermination(timeout.toMillis(), MILLISECONDS)).isTrue();
94         } catch (InterruptedException e) {
95           Thread.currentThread().interrupt();
96           fail("Error shutting down executor service:" + e);
97         } catch (Exception e) {
98           fail("Error shutting down executor service:" + e);
99         }
100       }
101 
102       if (!uncaughtExceptions.isEmpty()) {
103         String message =
104             uncaughtExceptions.stream()
105                 .map(e -> "\n\t" + getStackTraceAsString(e).replace("\t", "\t\t"))
106                 .collect(joining("\n"));
107         fail("Uncaught exceptions found: " + message);
108       }
109     } finally {
110       uncaughtExceptions.clear();
111       executorServices.clear();
112     }
113   }
114 }
115