1 /*
2  * Copyright (C) 2018 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.dx.mockito.inline.extended;
18 
19 import org.mockito.Mockito;
20 import org.mockito.MockitoSession;
21 import org.mockito.quality.Strictness;
22 
23 import java.util.ArrayList;
24 import java.util.HashMap;
25 
26 import static com.android.dx.mockito.inline.InlineStaticMockMaker.mockingInProgressClass;
27 
28 /**
29  * Same as {@link MockitoSession} but used when static methods are also stubbed.
30  */
31 @UnstableApi
32 public class StaticMockitoSession implements MockitoSession {
33     /**
34      * For each class where static mocking is enabled there is one marker object.
35      */
36     private static final HashMap<Class, Object> classToMarker = new HashMap<>();
37 
38     private final MockitoSession instanceSession;
39     private final ArrayList<Class<?>> staticMocks = new ArrayList<>(0);
40 
StaticMockitoSession(MockitoSession instanceSession)41     StaticMockitoSession(MockitoSession instanceSession) {
42         ExtendedMockito.addSession(this);
43         this.instanceSession = instanceSession;
44     }
45 
46     @Override
setStrictness(Strictness strictness)47     public void setStrictness(Strictness strictness) {
48         instanceSession.setStrictness(strictness);
49     }
50 
51     /**
52      * {@inheritDoc}
53      * <p><b>Extension:</b> This also resets all stubbing of static methods set up in the
54      * {@link ExtendedMockito#mockitoSession() builder} of the session.
55      */
56     @Override
finishMocking()57     public void finishMocking() {
58         finishMocking(null);
59     }
60 
61     /**
62      * {@inheritDoc}
63      * <p><b>Extension:</b> This also resets all stubbing of static methods set up in the
64      * {@link ExtendedMockito#mockitoSession() builder} of the session.
65      */
66     @Override
finishMocking(Throwable failure)67     public void finishMocking(Throwable failure) {
68         try {
69             instanceSession.finishMocking(failure);
70         } finally {
71             for (Class<?> clazz : staticMocks) {
72                 mockingInProgressClass.set(clazz);
73                 try {
74                     Mockito.reset(ExtendedMockito.staticMockMarker(clazz));
75                 } finally {
76                     mockingInProgressClass.remove();
77                 }
78                 classToMarker.remove(clazz);
79             }
80 
81             ExtendedMockito.removeSession(this);
82         }
83     }
84 
85     /**
86      * Init mocking for a class.
87      *
88      * @param mocking Description and settings of the mocking
89      * @param <T>     The class to mock
90      */
mockStatic(StaticMocking<T> mocking)91     <T> void mockStatic(StaticMocking<T> mocking) {
92         if (ExtendedMockito.staticMockMarker(mocking.clazz) != null) {
93             throw new IllegalArgumentException(mocking.clazz + " is already mocked");
94         }
95 
96         mockingInProgressClass.set(mocking.clazz);
97         try {
98             classToMarker.put(mocking.clazz, mocking.markerSupplier.get());
99         } finally {
100             mockingInProgressClass.remove();
101         }
102 
103         staticMocks.add(mocking.clazz);
104     }
105 
106     /**
107      * Get marker for a mocked/spies class or {@code null}.
108      *
109      * @param clazz The class that is mocked
110      * @return marker for a mocked class or {@code null} if class is not mocked in this session
111      * @see ExtendedMockito#staticMockMarker(Class)
112      */
113     @SuppressWarnings("unchecked")
staticMockMarker(Class<T> clazz)114     <T> T staticMockMarker(Class<T> clazz) {
115         return (T) classToMarker.get(clazz);
116     }
117 }
118