1 /*
2  * Copyright (C) 2021 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 com.android.server.tare;
18 
19 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
20 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
21 import static com.android.server.tare.TareTestUtils.assertLedgersEqual;
22 
23 import static org.junit.Assert.assertEquals;
24 import static org.mockito.ArgumentMatchers.anyInt;
25 import static org.mockito.ArgumentMatchers.anyString;
26 import static org.mockito.Mockito.mock;
27 
28 import android.app.AlarmManager;
29 import android.content.Context;
30 
31 import androidx.test.filters.SmallTest;
32 import androidx.test.runner.AndroidJUnit4;
33 
34 import com.android.server.LocalServices;
35 
36 import org.junit.After;
37 import org.junit.Before;
38 import org.junit.Test;
39 import org.junit.runner.RunWith;
40 import org.mockito.Mock;
41 import org.mockito.MockitoSession;
42 import org.mockito.quality.Strictness;
43 
44 /** Tests various aspects of the Agent. */
45 @RunWith(AndroidJUnit4.class)
46 @SmallTest
47 public class AgentTest {
48     private MockitoSession mMockingSession;
49     @Mock
50     private CompleteEconomicPolicy mEconomicPolicy;
51     @Mock
52     private Analyst mAnalyst;
53     @Mock
54     private Context mContext;
55     @Mock
56     private InternalResourceService mIrs;
57 
58     private Scribe mScribe;
59 
60     private static class MockScribe extends Scribe {
MockScribe(InternalResourceService irs, Analyst analyst)61         MockScribe(InternalResourceService irs, Analyst analyst) {
62             super(irs, analyst);
63         }
64 
65         @Override
postWrite()66         void postWrite() {
67             // Do nothing
68         }
69     }
70 
71     @Before
setUp()72     public void setUp() {
73         mMockingSession = mockitoSession()
74                 .initMocks(this)
75                 .strictness(Strictness.LENIENT)
76                 .mockStatic(LocalServices.class)
77                 .startMocking();
78         doReturn(mContext).when(mIrs).getContext();
79         doReturn(mEconomicPolicy).when(mIrs).getCompleteEconomicPolicyLocked();
80         doReturn(mIrs).when(mIrs).getLock();
81         doReturn(mock(AlarmManager.class)).when(mContext).getSystemService(Context.ALARM_SERVICE);
82         mScribe = new MockScribe(mIrs, mAnalyst);
83     }
84 
85     @After
tearDown()86     public void tearDown() {
87         if (mMockingSession != null) {
88             mMockingSession.finishMocking();
89         }
90     }
91 
92     @Test
testAppRemoval()93     public void testAppRemoval() {
94         final long consumptionLimit = 1_000_000L;
95         final long remainingCakes = consumptionLimit / 2;
96         mScribe.setConsumptionLimitLocked(consumptionLimit);
97         mScribe.adjustRemainingConsumableCakesLocked(remainingCakes - consumptionLimit);
98         assertEquals(remainingCakes, mScribe.getRemainingConsumableCakesLocked());
99 
100         final int userId = 0;
101         final String pkgName = "com.test";
102         final Agent agent = new Agent(mIrs, mScribe, mAnalyst);
103         final Ledger ledger = mScribe.getLedgerLocked(userId, pkgName);
104 
105         doReturn(consumptionLimit).when(mIrs).getConsumptionLimitLocked();
106         doReturn(consumptionLimit).when(mEconomicPolicy)
107                 .getMaxSatiatedBalance(anyInt(), anyString());
108 
109         Ledger.Transaction transaction = new Ledger.Transaction(0, 0, 0, null, 5, 10);
110         agent.recordTransactionLocked(userId, pkgName, ledger, transaction, false);
111         assertEquals(5, ledger.getCurrentBalance());
112         assertEquals(remainingCakes - 10, mScribe.getRemainingConsumableCakesLocked());
113 
114         agent.onPackageRemovedLocked(userId, pkgName);
115         assertEquals(remainingCakes - 10, mScribe.getRemainingConsumableCakesLocked());
116         assertLedgersEqual(new Ledger(), mScribe.getLedgerLocked(userId, pkgName));
117     }
118 
119     @Test
testRecordTransaction_UnderMax()120     public void testRecordTransaction_UnderMax() {
121         Agent agent = new Agent(mIrs, mScribe, mAnalyst);
122         Ledger ledger = new Ledger();
123 
124         doReturn(1_000_000L).when(mIrs).getConsumptionLimitLocked();
125         doReturn(1_000_000L).when(mEconomicPolicy).getMaxSatiatedBalance(anyInt(), anyString());
126 
127         Ledger.Transaction transaction = new Ledger.Transaction(0, 0, 0, null, 5, 0);
128         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
129         assertEquals(5, ledger.getCurrentBalance());
130 
131         transaction = new Ledger.Transaction(0, 0, 0, null, 995, 0);
132         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
133         assertEquals(1000, ledger.getCurrentBalance());
134 
135         transaction = new Ledger.Transaction(0, 0, 0, null, -500, 250);
136         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
137         assertEquals(500, ledger.getCurrentBalance());
138 
139         transaction = new Ledger.Transaction(0, 0, 0, null, 999_500L, 500);
140         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
141         assertEquals(1_000_000L, ledger.getCurrentBalance());
142 
143         transaction = new Ledger.Transaction(0, 0, 0, null, -1_000_001L, 1000);
144         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
145         assertEquals(-1, ledger.getCurrentBalance());
146     }
147 
148     @Test
testRecordTransaction_MaxConsumptionLimit()149     public void testRecordTransaction_MaxConsumptionLimit() {
150         Agent agent = new Agent(mIrs, mScribe, mAnalyst);
151         Ledger ledger = new Ledger();
152 
153         doReturn(1000L).when(mIrs).getConsumptionLimitLocked();
154         doReturn(1_000_000L).when(mEconomicPolicy).getMaxSatiatedBalance(anyInt(), anyString());
155 
156         Ledger.Transaction transaction = new Ledger.Transaction(0, 0, 0, null, 5, 0);
157         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
158         assertEquals(5, ledger.getCurrentBalance());
159 
160         transaction = new Ledger.Transaction(0, 0, 0, null, 995, 0);
161         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
162         assertEquals(1000, ledger.getCurrentBalance());
163 
164         transaction = new Ledger.Transaction(0, 0, 0, null, -500, 250);
165         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
166         assertEquals(500, ledger.getCurrentBalance());
167 
168         transaction = new Ledger.Transaction(0, 0, 0, null, 2000, 0);
169         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
170         assertEquals(2500, ledger.getCurrentBalance());
171 
172         // ConsumptionLimit can change as the battery level changes. Ledger balances shouldn't be
173         // affected.
174         doReturn(900L).when(mIrs).getConsumptionLimitLocked();
175 
176         transaction = new Ledger.Transaction(0, 0, 0, null, 100, 0);
177         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
178         assertEquals(2600, ledger.getCurrentBalance());
179 
180         transaction = new Ledger.Transaction(0, 0, 0, null, -50, 50);
181         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
182         assertEquals(2550, ledger.getCurrentBalance());
183 
184         transaction = new Ledger.Transaction(0, 0, 0, null, -200, 100);
185         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
186         assertEquals(2350, ledger.getCurrentBalance());
187 
188         doReturn(800L).when(mIrs).getConsumptionLimitLocked();
189 
190         transaction = new Ledger.Transaction(0, 0, 0, null, 100, 0);
191         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
192         assertEquals(2450, ledger.getCurrentBalance());
193     }
194 
195     @Test
testRecordTransaction_MaxSatiatedBalance()196     public void testRecordTransaction_MaxSatiatedBalance() {
197         Agent agent = new Agent(mIrs, mScribe, mAnalyst);
198         Ledger ledger = new Ledger();
199 
200         doReturn(1_000_000L).when(mIrs).getConsumptionLimitLocked();
201         doReturn(1000L).when(mEconomicPolicy).getMaxSatiatedBalance(anyInt(), anyString());
202 
203         Ledger.Transaction transaction = new Ledger.Transaction(0, 0, 0, null, 5, 0);
204         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
205         assertEquals(5, ledger.getCurrentBalance());
206 
207         transaction = new Ledger.Transaction(0, 0, 0, null, 995, 0);
208         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
209         assertEquals(1000, ledger.getCurrentBalance());
210 
211         transaction = new Ledger.Transaction(0, 0, 0, null, -500, 250);
212         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
213         assertEquals(500, ledger.getCurrentBalance());
214 
215         transaction = new Ledger.Transaction(0, 0, 0, null, 999_500L, 1000);
216         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
217         assertEquals(1_000, ledger.getCurrentBalance());
218 
219         // Shouldn't change in normal operation, but adding test case in case it does.
220         doReturn(900L).when(mEconomicPolicy).getMaxSatiatedBalance(anyInt(), anyString());
221 
222         transaction = new Ledger.Transaction(0, 0, 0, null, 500, 0);
223         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
224         assertEquals(1_000, ledger.getCurrentBalance());
225 
226         transaction = new Ledger.Transaction(0, 0, 0, null, -1001, 500);
227         agent.recordTransactionLocked(0, "com.test", ledger, transaction, false);
228         assertEquals(-1, ledger.getCurrentBalance());
229     }
230 }
231