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.common;
18 
19 import static com.google.common.truth.Truth.assertThat;
20 
21 import static org.junit.Assert.assertThrows;
22 
23 import android.os.Parcel;
24 
25 import androidx.test.ext.junit.runners.AndroidJUnit4;
26 
27 import com.google.protobuf.ByteString;
28 
29 import org.junit.Test;
30 import org.junit.runner.RunWith;
31 
32 import javax.annotation.Nullable;
33 
34 @RunWith(AndroidJUnit4.class)
35 public final class ExampleConsumptionTest {
36 
37     @Test
testBuilder_emptyTaskName()38     public void testBuilder_emptyTaskName() {
39         assertThrows(
40                 IllegalArgumentException.class,
41                 () ->
42                         new ExampleConsumption.Builder()
43                                 .setTaskId("")
44                                 .setExampleCount(10)
45                                 .setSelectionCriteria(new byte[] {10, 0, 1})
46                                 .build());
47     }
48 
49     @Test
testBuilder_nullTaskName()50     public void testBuilder_nullTaskName() {
51         assertThrows(
52                 IllegalArgumentException.class,
53                 () ->
54                         new ExampleConsumption.Builder()
55                                 .setTaskId(null)
56                                 .setExampleCount(10)
57                                 .setSelectionCriteria(new byte[] {10, 0, 1})
58                                 .build());
59     }
60 
61     @Test
testBuilder_normalCaseWithoutResumptionToken()62     public void testBuilder_normalCaseWithoutResumptionToken() {
63         String taskId = "my_task";
64         byte[] selectionCriteria = new byte[] {10, 0, 1};
65         int exampleCount = 10;
66         ExampleConsumption consumption =
67                 createExampleConsumption(taskId, selectionCriteria, exampleCount, null);
68         assertThat(consumption.getTaskId()).isEqualTo(taskId);
69         assertThat(consumption.getExampleCount()).isEqualTo(exampleCount);
70         assertThat(ByteString.copyFrom(consumption.getSelectionCriteria()))
71                 .isEqualTo(ByteString.copyFrom(selectionCriteria));
72         assertThat(consumption.getResumptionToken()).isNull();
73     }
74 
75     @Test
testBuilder_normalCaseWithResumptionToken()76     public void testBuilder_normalCaseWithResumptionToken() {
77         String taskId = "my_task";
78         byte[] selectionCriteria = new byte[] {10, 0, 1};
79         int exampleCount = 10;
80         byte[] resumptionToken = new byte[] {25, 10, 4, 56};
81         ExampleConsumption consumption =
82                 createExampleConsumption(taskId, selectionCriteria, exampleCount, resumptionToken);
83         assertThat(consumption.getTaskId()).isEqualTo(taskId);
84         assertThat(consumption.getExampleCount()).isEqualTo(exampleCount);
85         assertThat(ByteString.copyFrom(consumption.getSelectionCriteria()))
86                 .isEqualTo(ByteString.copyFrom(selectionCriteria));
87         assertThat(ByteString.copyFrom(consumption.getResumptionToken()))
88                 .isEqualTo(ByteString.copyFrom(resumptionToken));
89     }
90 
91     @Test
testWriteToParcel()92     public void testWriteToParcel() {
93         String taskId = "my_task";
94         byte[] selectionCriteria = new byte[] {10, 0, 1};
95         int exampleCount = 10;
96         byte[] resumptionToken = new byte[] {25, 10, 4, 56};
97         ExampleConsumption consumption =
98                 createExampleConsumption(taskId, selectionCriteria, exampleCount, resumptionToken);
99 
100         Parcel parcel = Parcel.obtain();
101         consumption.writeToParcel(parcel, 0);
102 
103         // Reset data position before recreating the parcelable.
104         parcel.setDataPosition(0);
105         ExampleConsumption recoveredConsumption =
106                 ExampleConsumption.CREATOR.createFromParcel(parcel);
107         assertThat(recoveredConsumption.getTaskId()).isEqualTo(taskId);
108         assertThat(recoveredConsumption.getExampleCount()).isEqualTo(exampleCount);
109         assertThat(ByteString.copyFrom(recoveredConsumption.getSelectionCriteria()))
110                 .isEqualTo(ByteString.copyFrom(selectionCriteria));
111         assertThat(ByteString.copyFrom(recoveredConsumption.getResumptionToken()))
112                 .isEqualTo(ByteString.copyFrom(resumptionToken));
113     }
114 
createExampleConsumption( String taskId, byte[] selectionCriteria, int exampleCount, @Nullable byte[] resumptionToken)115     private static ExampleConsumption createExampleConsumption(
116             String taskId,
117             byte[] selectionCriteria,
118             int exampleCount,
119             @Nullable byte[] resumptionToken) {
120         ExampleConsumption.Builder builder =
121                 new ExampleConsumption.Builder()
122                         .setTaskId(taskId)
123                         .setSelectionCriteria(selectionCriteria)
124                         .setExampleCount(exampleCount);
125         if (resumptionToken != null) {
126             builder.setResumptionToken(resumptionToken);
127         }
128         return builder.build();
129     }
130 }
131