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.federatedcompute;
18 
19 import static android.federatedcompute.common.ClientConstants.STATUS_INTERNAL_ERROR;
20 
21 import static com.google.common.truth.Truth.assertThat;
22 
23 import static org.junit.Assert.assertThrows;
24 import static org.mockito.ArgumentMatchers.any;
25 import static org.mockito.ArgumentMatchers.anyInt;
26 import static org.mockito.ArgumentMatchers.eq;
27 import static org.mockito.Mockito.doThrow;
28 import static org.mockito.Mockito.mock;
29 import static org.mockito.Mockito.never;
30 import static org.mockito.Mockito.times;
31 import static org.mockito.Mockito.verify;
32 
33 import android.federatedcompute.ExampleStoreQueryCallbackImpl.IteratorAdapter;
34 import android.federatedcompute.ExampleStoreQueryCallbackImpl.IteratorCallbackAdapter;
35 import android.federatedcompute.aidl.IExampleStoreCallback;
36 import android.federatedcompute.aidl.IExampleStoreIteratorCallback;
37 import android.os.Bundle;
38 import android.os.RemoteException;
39 import android.os.TransactionTooLargeException;
40 
41 import androidx.test.ext.junit.runners.AndroidJUnit4;
42 
43 import org.junit.Before;
44 import org.junit.Test;
45 import org.junit.runner.RunWith;
46 import org.mockito.ArgumentCaptor;
47 
48 @RunWith(AndroidJUnit4.class)
49 public final class ExampleStoreQueryCallbackImplTest {
50     private IteratorAdapter mMockIteratorAdaptor;
51     private ExampleStoreIterator mMockExampleStoreIterator;
52     private IExampleStoreIteratorCallback mMockAidlExampleStoreIteratorCallback;
53     private IExampleStoreCallback mMockAidlExampleStoreCallback;
54 
55     @Before
setUp()56     public void setUp() {
57         mMockIteratorAdaptor = mock(IteratorAdapter.class);
58         mMockExampleStoreIterator = mock(ExampleStoreIterator.class);
59         mMockAidlExampleStoreIteratorCallback = mock(IExampleStoreIteratorCallback.class);
60         mMockAidlExampleStoreCallback = mock(IExampleStoreCallback.class);
61     }
62 
63     @Test
testStartQuerySuccessNullResultThrows()64     public void testStartQuerySuccessNullResultThrows() throws Exception {
65         ExampleStoreQueryCallbackImpl adapter =
66                 new ExampleStoreQueryCallbackImpl(mMockAidlExampleStoreCallback);
67         assertThrows(NullPointerException.class, () -> adapter.onStartQuerySuccess(null));
68     }
69 
70     @Test
testStartQueryFailureTwicePassThrough()71     public void testStartQueryFailureTwicePassThrough() throws Exception {
72         ExampleStoreQueryCallbackImpl adapter =
73                 new ExampleStoreQueryCallbackImpl(mMockAidlExampleStoreCallback);
74         adapter.onStartQueryFailure(STATUS_INTERNAL_ERROR);
75         adapter.onStartQueryFailure(STATUS_INTERNAL_ERROR);
76         verify(mMockAidlExampleStoreCallback, times(2))
77                 .onStartQueryFailure(eq(STATUS_INTERNAL_ERROR));
78     }
79 
80     @Test
testStartQuerySuccessTwicePassThrough()81     public void testStartQuerySuccessTwicePassThrough() throws Exception {
82         ExampleStoreQueryCallbackImpl adapter =
83                 new ExampleStoreQueryCallbackImpl(mMockAidlExampleStoreCallback);
84         adapter.onStartQuerySuccess(mMockExampleStoreIterator);
85         adapter.onStartQuerySuccess(mMockExampleStoreIterator);
86         verify(mMockAidlExampleStoreCallback, times(2)).onStartQuerySuccess(any());
87     }
88 
89     @Test
testStartQuerySuccessAfterFailurePassThrough()90     public void testStartQuerySuccessAfterFailurePassThrough() throws Exception {
91         ExampleStoreQueryCallbackImpl adapter =
92                 new ExampleStoreQueryCallbackImpl(mMockAidlExampleStoreCallback);
93         adapter.onStartQueryFailure(STATUS_INTERNAL_ERROR);
94         adapter.onStartQuerySuccess(mMockExampleStoreIterator);
95         verify(mMockAidlExampleStoreCallback).onStartQuerySuccess(any());
96     }
97 
98     @Test
testStartQueryFailureAfterSuccessPassThrough()99     public void testStartQueryFailureAfterSuccessPassThrough() throws Exception {
100         ExampleStoreQueryCallbackImpl adapter =
101                 new ExampleStoreQueryCallbackImpl(mMockAidlExampleStoreCallback);
102         adapter.onStartQuerySuccess(mMockExampleStoreIterator);
103         adapter.onStartQueryFailure(STATUS_INTERNAL_ERROR);
104         verify(mMockAidlExampleStoreCallback).onStartQueryFailure(eq(STATUS_INTERNAL_ERROR));
105     }
106 
107     @Test
testIteratorCloseTwice()108     public void testIteratorCloseTwice() throws Exception {
109         ExampleStoreQueryCallbackImpl adapter =
110                 new ExampleStoreQueryCallbackImpl(mMockAidlExampleStoreCallback);
111         adapter.onStartQuerySuccess(mMockExampleStoreIterator);
112         ArgumentCaptor<IteratorAdapter> iteratorCaptor =
113                 ArgumentCaptor.forClass(IteratorAdapter.class);
114         verify(mMockAidlExampleStoreCallback).onStartQuerySuccess(iteratorCaptor.capture());
115         IteratorAdapter iterator = iteratorCaptor.getValue();
116         iterator.close();
117         iterator.close();
118         // The app's iterator should only have been called once.
119         verify(mMockExampleStoreIterator).close();
120     }
121 
122     @Test
testIteratorAdapterCloseTwiceIgnoresSecondCall()123     public void testIteratorAdapterCloseTwiceIgnoresSecondCall() throws Exception {
124         IteratorAdapter iterator = new IteratorAdapter(mMockExampleStoreIterator);
125         iterator.close();
126         verify(mMockExampleStoreIterator).close();
127         iterator.close();
128         // The second call shouldn't result in another call to the app's close() method.
129         verify(mMockExampleStoreIterator).close();
130     }
131 
132     @Test
testIteratorAdapterNextAfterCloseIgnore()133     public void testIteratorAdapterNextAfterCloseIgnore() throws Exception {
134         IteratorAdapter iterator = new IteratorAdapter(mMockExampleStoreIterator);
135         iterator.close();
136         iterator.next(mMockAidlExampleStoreIteratorCallback);
137         // The second call shouldn't result in another call to the app's close() method.
138         verify(mMockExampleStoreIterator, never()).next(any());
139     }
140 
141     /**
142      * Tests that additional calls to a the callback are passed through to the proxy. It will be in
143      * charge of ignoring all but the first call.
144      */
145     @Test
testIteratorCallbackSuccessTwicePassThrough()146     public void testIteratorCallbackSuccessTwicePassThrough() throws Exception {
147         IteratorCallbackAdapter adapter =
148                 new IteratorCallbackAdapter(
149                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
150         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isTrue();
151         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isTrue();
152         verify(mMockAidlExampleStoreIteratorCallback, times(2)).onIteratorNextSuccess(any());
153     }
154 
155     /**
156      * Tests that additional calls to a the callback are passed through to the proxy. It will be in
157      * charge of ignoring all but the first call.
158      */
159     @Test
testIteratorCallbackFailureTwicePassThrough()160     public void testIteratorCallbackFailureTwicePassThrough() throws Exception {
161         IteratorCallbackAdapter adapter =
162                 new IteratorCallbackAdapter(
163                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
164         adapter.onIteratorNextFailure(STATUS_INTERNAL_ERROR);
165         adapter.onIteratorNextFailure(STATUS_INTERNAL_ERROR);
166         verify(mMockAidlExampleStoreIteratorCallback, times(2))
167                 .onIteratorNextFailure(eq(STATUS_INTERNAL_ERROR));
168     }
169 
170     /**
171      * Tests that additional calls to a the callback are passed through to the proxy. It will be in
172      * charge of ignoring all but the first call.
173      */
174     @Test
testIteratorCallbackSuccessAfterFailurePassThrough()175     public void testIteratorCallbackSuccessAfterFailurePassThrough() throws Exception {
176         IteratorCallbackAdapter adapter =
177                 new IteratorCallbackAdapter(
178                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
179         adapter.onIteratorNextFailure(STATUS_INTERNAL_ERROR);
180         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isTrue();
181         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isTrue();
182         verify(mMockAidlExampleStoreIteratorCallback)
183                 .onIteratorNextFailure(eq(STATUS_INTERNAL_ERROR));
184         verify(mMockAidlExampleStoreIteratorCallback, times(2)).onIteratorNextSuccess(any());
185     }
186 
187     /**
188      * Tests that additional calls to a the callback are passed through to the proxy. It will be in
189      * charge of ignoring all but the first call.
190      */
191     @Test
testIteratorCallbackFailureAfterSuccessPassThrough()192     public void testIteratorCallbackFailureAfterSuccessPassThrough() throws Exception {
193         IteratorCallbackAdapter adapter =
194                 new IteratorCallbackAdapter(
195                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
196         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isTrue();
197         adapter.onIteratorNextFailure(STATUS_INTERNAL_ERROR);
198         verify(mMockAidlExampleStoreIteratorCallback).onIteratorNextSuccess(any());
199         verify(mMockAidlExampleStoreIteratorCallback)
200                 .onIteratorNextFailure(eq(STATUS_INTERNAL_ERROR));
201     }
202 
203     @Test
testIteratorCallbackSuccessRemoteException()204     public void testIteratorCallbackSuccessRemoteException() throws Exception {
205         doThrow(new RemoteException())
206                 .when(mMockAidlExampleStoreIteratorCallback)
207                 .onIteratorNextSuccess(any());
208         IteratorCallbackAdapter adapter =
209                 new IteratorCallbackAdapter(
210                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
211         // Should not throw. Exception should be swallowed, but an error should be returned.
212         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isFalse();
213         // The corresponding iterator should also be closed when a RemoteException occurs.
214         verify(mMockIteratorAdaptor).close();
215     }
216 
217     @Test
testIteratorCallbackSuccessTransactionTooLargeException()218     public void testIteratorCallbackSuccessTransactionTooLargeException() throws Exception {
219         doThrow(new TransactionTooLargeException())
220                 .when(mMockAidlExampleStoreIteratorCallback)
221                 .onIteratorNextSuccess(any());
222         IteratorCallbackAdapter adapter =
223                 new IteratorCallbackAdapter(
224                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
225         // Should return an error to the client so they can at least debug when examples are
226         // dropped.
227         assertThat(adapter.onIteratorNextSuccess(new Bundle())).isFalse();
228         // The corresponding iterator should be closed automatically in this case.
229         verify(mMockIteratorAdaptor).close();
230     }
231 
232     @Test
testIteratorCallbackFailureRemoteException()233     public void testIteratorCallbackFailureRemoteException() throws Exception {
234         doThrow(new RemoteException())
235                 .when(mMockAidlExampleStoreIteratorCallback)
236                 .onIteratorNextFailure(anyInt());
237         IteratorCallbackAdapter adapter =
238                 new IteratorCallbackAdapter(
239                         mMockAidlExampleStoreIteratorCallback, mMockIteratorAdaptor);
240         // Should not throw. Exception should be swallowed.
241         adapter.onIteratorNextFailure(STATUS_INTERNAL_ERROR);
242         // The corresponding iterator should also be closed when a RemoteException occurs.
243         verify(mMockIteratorAdaptor).close();
244     }
245 }
246