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 com.android.adservices.shared.testing.junit;
18 
19 import org.junit.Rule;
20 import org.junit.rules.MethodRule;
21 import org.junit.rules.TestRule;
22 import org.junit.runner.Description;
23 import org.junit.runners.model.FrameworkMethod;
24 import org.junit.runners.model.Statement;
25 
26 import java.util.ArrayList;
27 import java.util.Collections;
28 import java.util.Comparator;
29 import java.util.IdentityHashMap;
30 import java.util.List;
31 
32 /**
33  * Data structure for ordering of {@link TestRule}/{@link MethodRule} instances. Copied from package
34  * {@link org.junit.runners} which has RuleContainer class as this is package protected.
35  */
36 public final class RuleContainer {
37     private final IdentityHashMap<Object, Integer> mOrderValues =
38             new IdentityHashMap<Object, Integer>();
39     private final List<TestRule> testRules = new ArrayList<TestRule>();
40     private final List<MethodRule> methodRules = new ArrayList<MethodRule>();
41 
42     /** Sets order value for the specified rule. */
setOrder(Object rule, int order)43     public void setOrder(Object rule, int order) {
44         mOrderValues.put(rule, order);
45     }
46 
add(MethodRule methodRule)47     public void add(MethodRule methodRule) {
48         methodRules.add(methodRule);
49     }
50 
add(TestRule testRule)51     public void add(TestRule testRule) {
52         testRules.add(testRule);
53     }
54 
55     static final Comparator<RuleEntry> ENTRY_COMPARATOR =
56             new Comparator<RuleEntry>() {
57                 public int compare(RuleEntry o1, RuleEntry o2) {
58                     int result = compareInt(o1.order, o2.order);
59                     return result != 0 ? result : o1.type - o2.type;
60                 }
61 
62                 private int compareInt(int a, int b) {
63                     return (a < b) ? 1 : (a == b ? 0 : -1);
64                 }
65             };
66 
67     /** Returns entries in the order how they should be applied, i.e. inner-to-outer. */
getSortedEntries()68     private List<RuleEntry> getSortedEntries() {
69         List<RuleEntry> ruleEntries =
70                 new ArrayList<RuleEntry>(methodRules.size() + testRules.size());
71         for (MethodRule rule : methodRules) {
72             ruleEntries.add(
73                     new RuleEntry(rule, RuleEntry.TYPE_METHOD_RULE, mOrderValues.get(rule)));
74         }
75         for (TestRule rule : testRules) {
76             ruleEntries.add(new RuleEntry(rule, RuleEntry.TYPE_TEST_RULE, mOrderValues.get(rule)));
77         }
78         Collections.sort(ruleEntries, ENTRY_COMPARATOR);
79         return ruleEntries;
80     }
81 
82     /** Applies all the rules ordered accordingly to the specified {@code statement}. */
apply( FrameworkMethod method, Description description, Object target, Statement statement)83     public Statement apply(
84             FrameworkMethod method, Description description, Object target, Statement statement) {
85         if (methodRules.isEmpty() && testRules.isEmpty()) {
86             return statement;
87         }
88         Statement result = statement;
89         for (RuleEntry ruleEntry : getSortedEntries()) {
90             if (ruleEntry.type == RuleEntry.TYPE_TEST_RULE) {
91                 result = ((TestRule) ruleEntry.rule).apply(result, description);
92             } else {
93                 result = ((MethodRule) ruleEntry.rule).apply(result, method, target);
94             }
95         }
96         return result;
97     }
98 
99     /**
100      * Returns rule instances in the order how they should be applied, i.e. inner-to-outer.
101      * VisibleForTesting
102      */
getSortedRules()103     List<Object> getSortedRules() {
104         List<Object> result = new ArrayList<Object>();
105         for (RuleEntry entry : getSortedEntries()) {
106             result.add(entry.rule);
107         }
108         return result;
109     }
110 
111     static class RuleEntry {
112         static final int TYPE_TEST_RULE = 1;
113         static final int TYPE_METHOD_RULE = 0;
114 
115         final Object rule;
116         final int type;
117         final int order;
118 
RuleEntry(Object rule, int type, Integer order)119         RuleEntry(Object rule, int type, Integer order) {
120             this.rule = rule;
121             this.type = type;
122             this.order = order != null ? order.intValue() : Rule.DEFAULT_ORDER;
123         }
124     }
125 }
126