1 /*
2  * Copyright (C) 2019 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 package com.android.customization.testutils;
17 
18 import static org.mockito.ArgumentMatchers.any;
19 import static org.mockito.ArgumentMatchers.anyInt;
20 import static org.mockito.ArgumentMatchers.anyString;
21 import static org.mockito.Mockito.when;
22 
23 import android.content.om.OverlayManager;
24 import android.text.TextUtils;
25 
26 import com.android.customization.model.theme.OverlayManagerCompat;
27 
28 import org.mockito.stubbing.Answer;
29 
30 import java.util.Arrays;
31 import java.util.HashSet;
32 import java.util.Map;
33 import java.util.Set;
34 import java.util.function.Function;
35 import java.util.stream.Collectors;
36 
37 /**
38  * Helper class to provide mock implementation for OverlayManager, to use, create a Mockito Mock
39  * for OverlayManager and call {@link #setUpMock(OverlayManager)} with it, then use
40  * {@link #addOverlay(String, String, String, boolean, int)} to add fake OverlayInfo to be returned
41  * by the mocked OverlayManager.
42  */
43 public class OverlayManagerMocks {
44     private static class MockOverlay {
45         final String packageName;
46         final String targetPackage;
47         final String category;
48 
MockOverlay(String packageName, String targetPackage, String category)49         public MockOverlay(String packageName, String targetPackage, String category) {
50             this.packageName = packageName;
51             this.targetPackage = targetPackage;
52             this.category = category;
53         }
54 
55         @Override
equals(Object obj)56         public boolean equals(Object obj) {
57             return obj instanceof MockOverlay
58                     && TextUtils.equals(((MockOverlay) obj).packageName, packageName)
59                     && TextUtils.equals(((MockOverlay) obj).targetPackage, targetPackage)
60                     && TextUtils.equals(((MockOverlay) obj).category, category);
61         }
62     }
63 
64     private Set<MockOverlay> mAllOverlays = new HashSet<>();
65     private Set<MockOverlay> mEnabledOverlays = new HashSet<>();
66 
setEnabled(String packageName, boolean enable, int userId)67     private boolean setEnabled(String packageName, boolean enable, int userId) {
68         if (packageName == null) {
69             return false;
70         }
71         Set<MockOverlay> packageOverlays = mAllOverlays.stream()
72                 .filter(mockOverlay -> mockOverlay.packageName.equals(packageName)).collect(
73                 Collectors.toSet());;
74         if (packageOverlays.isEmpty()) {
75             return false;
76         }
77         if (enable) {
78             mEnabledOverlays.addAll(packageOverlays);
79         } else {
80             mEnabledOverlays.removeAll(packageOverlays);
81         }
82         return true;
83     }
84 
addOverlay(String packageName, String targetPackage, String category, boolean enabled, int userId)85     public void addOverlay(String packageName, String targetPackage, String category,
86             boolean enabled, int userId) {
87         MockOverlay overlay = new MockOverlay(packageName, targetPackage, category);
88         mAllOverlays.add(overlay);
89         if (enabled) {
90             mEnabledOverlays.add(overlay);
91         }
92     }
93 
clearOverlays()94     public void clearOverlays() {
95         mAllOverlays.clear();
96         mEnabledOverlays.clear();
97     }
98 
setUpMock(OverlayManagerCompat mockOverlayManager)99     public void setUpMock(OverlayManagerCompat mockOverlayManager) {
100         when(mockOverlayManager.getEnabledPackageName(anyString(), anyString())).then(
101                 (Answer<String>) inv ->
102                         mEnabledOverlays.stream().filter(
103                                 mockOverlay ->
104                                         mockOverlay.targetPackage.equals(inv.getArgument(0))
105                                             && mockOverlay.category.equals(inv.getArgument(1)))
106                                 .map(overlay -> overlay.packageName).findFirst().orElse(null));
107 
108 
109         when(mockOverlayManager.disableOverlay(anyString(), anyInt())).then(
110                 (Answer<Boolean>) invocation ->
111                         setEnabled(invocation.getArgument(0),
112                                 false,
113                                 invocation.getArgument(1)));
114 
115         when(mockOverlayManager.setEnabledExclusiveInCategory(anyString(), anyInt())).then(
116                 (Answer<Boolean>) invocation ->
117                         setEnabled(
118                                 invocation.getArgument(0),
119                                 true,
120                                 invocation.getArgument(1)));
121 
122         when(mockOverlayManager.getEnabledOverlaysForTargets(any())).then(
123                 (Answer<Map<String, String>>) inv ->
124                         mEnabledOverlays.stream().filter(
125                                 overlay ->
126                                         Arrays.asList(inv.getArguments())
127                                                 .contains(overlay.targetPackage))
128                                 .collect(Collectors.toMap(
129                                         overlay ->
130                                                 overlay.category,
131                                         (Function<MockOverlay, String>) overlay ->
132                                                 overlay.packageName))
133         );
134     }
135 }