m-trieu commented on code in PR #30572:
URL: https://github.com/apache/beam/pull/30572#discussion_r1518189275
##########
runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/budget/GetWorkBudgetRefresherTest.java:
##########
@@ -17,89 +17,105 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill.work.budget;
+import static com.google.common.truth.Truth.assertThat;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.Timeout;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
-import org.mockito.Mockito;
@RunWith(JUnit4.class)
public class GetWorkBudgetRefresherTest {
- @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
private static final int WAIT_BUFFER = 10;
- private final Runnable redistributeBudget = Mockito.mock(Runnable.class);
+ @Rule public transient Timeout globalTimeout = Timeout.seconds(600);
- private GetWorkBudgetRefresher createBudgetRefresher() {
- return createBudgetRefresher(false);
+ private GetWorkBudgetRefresher createBudgetRefresher(Runnable
redistributeBudget) {
+ return createBudgetRefresher(false, redistributeBudget);
}
- private GetWorkBudgetRefresher createBudgetRefresher(Boolean
isBudgetRefreshPaused) {
+ private GetWorkBudgetRefresher createBudgetRefresher(
+ boolean isBudgetRefreshPaused, Runnable redistributeBudget) {
return new GetWorkBudgetRefresher(() -> isBudgetRefreshPaused,
redistributeBudget);
}
@Test
public void testStop_successfullyTerminates() throws InterruptedException {
- GetWorkBudgetRefresher budgetRefresher = createBudgetRefresher();
+ CountDownLatch redistributeBudgetLatch = new CountDownLatch(1);
+ Runnable redistributeBudget = redistributeBudgetLatch::countDown;
+ GetWorkBudgetRefresher budgetRefresher =
createBudgetRefresher(redistributeBudget);
budgetRefresher.start();
budgetRefresher.stop();
budgetRefresher.requestBudgetRefresh();
- Thread.sleep(WAIT_BUFFER);
+ redistributeBudgetLatch.await(WAIT_BUFFER, TimeUnit.MILLISECONDS);
+ // Make sure that redistributeBudgetLatch.countDown() is never called.
verifyNoInteractions(redistributeBudget);
+ assertThat(redistributeBudgetLatch.getCount()).isEqualTo(1);
}
@Test
public void testRequestBudgetRefresh_triggersBudgetRefresh() throws
InterruptedException {
- GetWorkBudgetRefresher budgetRefresher = createBudgetRefresher();
+ CountDownLatch redistributeBudgetLatch = new CountDownLatch(1);
+ Runnable redistributeBudget = redistributeBudgetLatch::countDown;
+ GetWorkBudgetRefresher budgetRefresher =
createBudgetRefresher(redistributeBudget);
budgetRefresher.start();
budgetRefresher.requestBudgetRefresh();
// Wait a bit for redistribute budget to run.
- Thread.sleep(WAIT_BUFFER);
+ redistributeBudgetLatch.await();
verify(redistributeBudget, times(1)).run();
}
@Test
public void testScheduledBudgetRefresh() throws InterruptedException {
- GetWorkBudgetRefresher budgetRefresher = createBudgetRefresher();
+ CountDownLatch redistributeBudgetLatch = new CountDownLatch(1);
+ Runnable redistributeBudget = redistributeBudgetLatch::countDown;
+ GetWorkBudgetRefresher budgetRefresher =
createBudgetRefresher(redistributeBudget);
budgetRefresher.start();
- Thread.sleep(GetWorkBudgetRefresher.SCHEDULED_BUDGET_REFRESH_MILLIS +
WAIT_BUFFER);
+ redistributeBudgetLatch.await();
verify(redistributeBudget, times(1)).run();
Review Comment:
removed,
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
