1 /*
2  * Copyright 2022 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.adservices.ondevicepersonalization;
18 
19 import static com.google.common.truth.Truth.assertThat;
20 
21 import static org.junit.Assert.assertThrows;
22 
23 import android.adservices.ondevicepersonalization.aidl.IDataAccessService;
24 import android.adservices.ondevicepersonalization.aidl.IDataAccessServiceCallback;
25 import android.adservices.ondevicepersonalization.aidl.IFederatedComputeCallback;
26 import android.adservices.ondevicepersonalization.aidl.IFederatedComputeService;
27 import android.federatedcompute.common.TrainingOptions;
28 import android.os.Bundle;
29 import android.os.RemoteException;
30 
31 import androidx.test.ext.junit.runners.AndroidJUnit4;
32 import androidx.test.filters.SmallTest;
33 
34 import org.junit.Test;
35 import org.junit.runner.RunWith;
36 
37 import java.time.Duration;
38 
39 /** Unit Tests of RemoteData API. */
40 @SmallTest
41 @RunWith(AndroidJUnit4.class)
42 public class FederatedComputeSchedulerTest {
43     FederatedComputeScheduler mFederatedComputeScheduler =
44             new FederatedComputeScheduler(
45                     IFederatedComputeService.Stub.asInterface(new FederatedComputeService()),
46                     IDataAccessService.Stub.asInterface(new TestDataService()));
47 
48     private boolean mCancelCalled = false;
49     private boolean mScheduleCalled = false;
50     private boolean mLogApiCalled = false;
51 
52     @Test
testScheduleSuccess()53     public void testScheduleSuccess() {
54         TrainingInterval interval =
55                 new TrainingInterval.Builder()
56                         .setMinimumInterval(Duration.ofHours(10))
57                         .setSchedulingMode(TrainingInterval.SCHEDULING_MODE_ONE_TIME)
58                         .build();
59         FederatedComputeScheduler.Params params = new FederatedComputeScheduler.Params(interval);
60         FederatedComputeInput input =
61                 new FederatedComputeInput.Builder().setPopulationName("population").build();
62         mFederatedComputeScheduler.schedule(params, input);
63         assertThat(mScheduleCalled).isTrue();
64         assertThat(mLogApiCalled).isTrue();
65     }
66 
67     @Test
testScheduleNull()68     public void testScheduleNull() {
69         FederatedComputeScheduler fcs = new FederatedComputeScheduler(null, new TestDataService());
70         TrainingInterval interval =
71                 new TrainingInterval.Builder()
72                         .setMinimumInterval(Duration.ofHours(10))
73                         .setSchedulingMode(TrainingInterval.SCHEDULING_MODE_ONE_TIME)
74                         .build();
75         FederatedComputeScheduler.Params params = new FederatedComputeScheduler.Params(interval);
76         FederatedComputeInput input =
77                 new FederatedComputeInput.Builder().setPopulationName("population").build();
78         assertThrows(IllegalStateException.class, () -> fcs.schedule(params, input));
79     }
80 
81     @Test
testScheduleErr()82     public void testScheduleErr() {
83         TrainingInterval interval =
84                 new TrainingInterval.Builder()
85                         .setMinimumInterval(Duration.ofHours(10))
86                         .setSchedulingMode(TrainingInterval.SCHEDULING_MODE_ONE_TIME)
87                         .build();
88         FederatedComputeScheduler.Params params = new FederatedComputeScheduler.Params(interval);
89         FederatedComputeInput input =
90                 new FederatedComputeInput.Builder().setPopulationName("err").build();
91         mFederatedComputeScheduler.schedule(params, input);
92         assertThat(mScheduleCalled).isTrue();
93         assertThat(mLogApiCalled).isTrue();
94     }
95 
96     @Test
testCancelSuccess()97     public void testCancelSuccess() {
98         FederatedComputeInput input =
99                 new FederatedComputeInput.Builder().setPopulationName("population").build();
100         mFederatedComputeScheduler.cancel(input);
101         assertThat(mCancelCalled).isTrue();
102         assertThat(mLogApiCalled).isTrue();
103     }
104 
105     @Test
testCancelNull()106     public void testCancelNull() {
107         FederatedComputeInput input =
108                 new FederatedComputeInput.Builder().setPopulationName("population").build();
109         FederatedComputeScheduler fcs = new FederatedComputeScheduler(null, new TestDataService());
110         assertThrows(IllegalStateException.class, () -> fcs.cancel(input));
111     }
112 
113     @Test
testCancelErr()114     public void testCancelErr() {
115         FederatedComputeInput input =
116                 new FederatedComputeInput.Builder().setPopulationName("err").build();
117         mFederatedComputeScheduler.cancel(input);
118         assertThat(mCancelCalled).isTrue();
119         assertThat(mLogApiCalled).isTrue();
120     }
121 
122     class FederatedComputeService extends IFederatedComputeService.Stub {
123         @Override
schedule( TrainingOptions trainingOptions, IFederatedComputeCallback iFederatedComputeCallback)124         public void schedule(
125                 TrainingOptions trainingOptions,
126                 IFederatedComputeCallback iFederatedComputeCallback)
127                 throws RemoteException {
128             mScheduleCalled = true;
129             if (trainingOptions.getPopulationName().equals("err")) {
130                 iFederatedComputeCallback.onFailure(1);
131             }
132             iFederatedComputeCallback.onSuccess();
133         }
134 
135         @Override
cancel(String s, IFederatedComputeCallback iFederatedComputeCallback)136         public void cancel(String s, IFederatedComputeCallback iFederatedComputeCallback)
137                 throws RemoteException {
138             mCancelCalled = true;
139             if (s.equals("err")) {
140                 iFederatedComputeCallback.onFailure(1);
141             }
142             iFederatedComputeCallback.onSuccess();
143         }
144     }
145 
146     class TestDataService extends IDataAccessService.Stub {
147         @Override
onRequest(int operation, Bundle params, IDataAccessServiceCallback callback)148         public void onRequest(int operation, Bundle params, IDataAccessServiceCallback callback) {}
149 
150         @Override
logApiCallStats(int apiName, long latencyMillis, int responseCode)151         public void logApiCallStats(int apiName, long latencyMillis, int responseCode) {
152             mLogApiCalled = true;
153         }
154     }
155 }
156