1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.server.wm.jetpack.extensions.util;
18 
19 import androidx.annotation.NonNull;
20 import androidx.annotation.Nullable;
21 import androidx.window.extensions.core.util.function.Consumer;
22 
23 import java.util.ArrayList;
24 import java.util.Collections;
25 import java.util.List;
26 import java.util.concurrent.LinkedBlockingQueue;
27 import java.util.concurrent.TimeUnit;
28 
29 /**
30  * Consumer that provides a simple way to wait for a specific count of values to be received within
31  * a timeout and then return the last value.
32  *
33  * It uses extensions core version of {@link Consumer} instead of
34  * {@link java.util.function.Consumer Java 8 version Consumer}.
35  */
36 public class TestValueCountConsumer<T> implements Consumer<T> {
37 
38     private static final long TIMEOUT_MS = 3000;
39     private static final int DEFAULT_COUNT = 1;
40     private int mCount = DEFAULT_COUNT;
41     private LinkedBlockingQueue<T> mLinkedBlockingQueue;
42     private T mLastReportedValue;
43 
TestValueCountConsumer()44     public TestValueCountConsumer() {
45         mLinkedBlockingQueue = new LinkedBlockingQueue<>();
46     }
47 
48     @Override
accept(T value)49     public void accept(T value) {
50         // Asynchronously offer value to queue
51         mLinkedBlockingQueue.offer(value);
52     }
53 
setCount(int count)54     public void setCount(int count) {
55         mCount = count;
56     }
57 
58     /**
59      * Returns the value that was reported after the count was reached from
60      * {@link TestValueCountConsumer#setCount(int)}.
61      */
62     @Nullable
waitAndGet()63     public T waitAndGet() throws InterruptedException {
64         T value = null;
65         for (int i = 0; i < mCount; i++) {
66             value = mLinkedBlockingQueue.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS);
67         }
68         mLastReportedValue = value;
69         return value;
70     }
71 
72     /**
73      * Returns a list that contains the number of values set in
74      * {@link TestValueCountConsumer#setCount(int)}.
75      */
76     @NonNull
waitAndGetAllValues()77     public List<T> waitAndGetAllValues() throws InterruptedException {
78         List<T> values = new ArrayList<>();
79         for (int i = 0; i < mCount; i++) {
80             T value = mLinkedBlockingQueue.poll(TIMEOUT_MS, TimeUnit.MILLISECONDS);
81             if (value != null) {
82                 values.add(value);
83             }
84         }
85         return Collections.unmodifiableList(values);
86     }
87 
88     /**
89      * Clears the queue of currently recorded values.
90      */
clearQueue()91     public void clearQueue() {
92         mLinkedBlockingQueue.clear();
93     }
94 
95     @Nullable
getLastReportedValue()96     public T getLastReportedValue() {
97         return mLastReportedValue;
98     }
99 }
100