Skip to content

Commit e8be56d

Browse files
lhotarilhotari
authored
[fix][test] Fix UnfinishedStubbing issue in AuthZTests (#24165)
1 parent 6e8c349 commit e8be56d

3 files changed

Lines changed: 70 additions & 48 deletions

File tree

‎pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/AuthZTest.java‎

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,22 +18,25 @@
1818
*/
1919
package org.apache.pulsar.broker.admin;
2020

21+
import static org.mockito.Mockito.doReturn;
2122
import io.jsonwebtoken.Jwts;
23+
import java.util.UUID;
24+
import java.util.concurrent.atomic.AtomicBoolean;
25+
import java.util.function.Consumer;
2226
import org.apache.commons.lang3.reflect.FieldUtils;
27+
import org.apache.pulsar.broker.BrokerTestUtil;
2328
import org.apache.pulsar.broker.authorization.AuthorizationService;
2429
import org.apache.pulsar.client.admin.PulsarAdmin;
2530
import org.apache.pulsar.common.policies.data.NamespaceOperation;
2631
import org.apache.pulsar.common.policies.data.TopicOperation;
2732
import org.apache.pulsar.security.MockedPulsarStandalone;
2833
import org.mockito.Mockito;
34+
import org.mockito.invocation.InvocationOnMock;
2935
import org.testng.Assert;
3036
import org.testng.annotations.AfterMethod;
3137
import org.testng.annotations.BeforeMethod;
32-
import java.util.UUID;
33-
import java.util.concurrent.atomic.AtomicBoolean;
34-
import static org.mockito.Mockito.doReturn;
3538

36-
public class AuthZTest extends MockedPulsarStandalone {
39+
public abstract class AuthZTest extends MockedPulsarStandalone {
3740

3841
protected PulsarAdmin superUserAdmin;
3942

@@ -47,6 +50,9 @@ public class AuthZTest extends MockedPulsarStandalone {
4750
protected static final String TENANT_ADMIN_TOKEN = Jwts.builder()
4851
.claim("sub", TENANT_ADMIN_SUBJECT).signWith(SECRET_KEY).compact();
4952

53+
private volatile Consumer<InvocationOnMock> allowTopicOperationAsyncHandler;
54+
private volatile Consumer<InvocationOnMock> allowNamespaceOperationAsyncHandler;
55+
5056
@Override
5157
public void close() throws Exception {
5258
if (superUserAdmin != null) {
@@ -65,48 +71,62 @@ public void close() throws Exception {
6571
@BeforeMethod(alwaysRun = true)
6672
public void before() throws IllegalAccessException {
6773
orignalAuthorizationService = getPulsarService().getBrokerService().getAuthorizationService();
68-
authorizationService = Mockito.spy(orignalAuthorizationService);
74+
authorizationService = BrokerTestUtil.spyWithoutRecordingInvocations(orignalAuthorizationService);
6975
FieldUtils.writeField(getPulsarService().getBrokerService(), "authorizationService",
7076
authorizationService, true);
77+
Mockito.doAnswer(invocationOnMock -> {
78+
Consumer<InvocationOnMock> localAllowTopicOperationAsyncHandler =
79+
allowTopicOperationAsyncHandler;
80+
if (localAllowTopicOperationAsyncHandler != null) {
81+
localAllowTopicOperationAsyncHandler.accept(invocationOnMock);
82+
}
83+
return invocationOnMock.callRealMethod();
84+
}).when(authorizationService).allowTopicOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
85+
Mockito.any(), Mockito.any());
86+
doReturn(true)
87+
.when(authorizationService).isValidOriginalPrincipal(Mockito.any(), Mockito.any(), Mockito.any());
88+
Mockito.doAnswer(invocationOnMock -> {
89+
Consumer<InvocationOnMock> localAllowNamespaceOperationAsyncHandler =
90+
allowNamespaceOperationAsyncHandler;
91+
if (localAllowNamespaceOperationAsyncHandler != null) {
92+
localAllowNamespaceOperationAsyncHandler.accept(invocationOnMock);
93+
}
94+
return invocationOnMock.callRealMethod();
95+
}).when(authorizationService).allowNamespaceOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
96+
Mockito.any(), Mockito.any());
7197
}
7298

7399
@AfterMethod(alwaysRun = true)
74100
public void after() throws IllegalAccessException {
75101
FieldUtils.writeField(getPulsarService().getBrokerService(), "authorizationService",
76102
orignalAuthorizationService, true);
103+
allowNamespaceOperationAsyncHandler = null;
104+
allowTopicOperationAsyncHandler = null;
77105
}
78106

79107
protected AtomicBoolean setAuthorizationTopicOperationChecker(String role, Object operation) {
80108
AtomicBoolean execFlag = new AtomicBoolean(false);
81109
if (operation instanceof TopicOperation) {
82-
Mockito.doAnswer(invocationOnMock -> {
110+
allowTopicOperationAsyncHandler = invocationOnMock -> {
83111
String role_ = invocationOnMock.getArgument(2);
84112
if (role.equals(role_)) {
85113
TopicOperation operation_ = invocationOnMock.getArgument(1);
86114
Assert.assertEquals(operation_, operation);
87115
}
88116
execFlag.set(true);
89-
return invocationOnMock.callRealMethod();
90-
}).when(authorizationService).allowTopicOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
91-
Mockito.any(), Mockito.any());
117+
};
92118
} else if (operation instanceof NamespaceOperation) {
93-
doReturn(true)
94-
.when(authorizationService).isValidOriginalPrincipal(Mockito.any(), Mockito.any(), Mockito.any());
95-
Mockito.doAnswer(invocationOnMock -> {
119+
allowNamespaceOperationAsyncHandler = invocationOnMock -> {
96120
String role_ = invocationOnMock.getArgument(2);
97121
if (role.equals(role_)) {
98122
TopicOperation operation_ = invocationOnMock.getArgument(1);
99123
Assert.assertEquals(operation_, operation);
100124
}
101125
execFlag.set(true);
102-
return invocationOnMock.callRealMethod();
103-
}).when(authorizationService).allowNamespaceOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
104-
Mockito.any(), Mockito.any());
126+
};
105127
} else {
106128
throw new IllegalArgumentException("");
107129
}
108-
109-
110130
return execFlag;
111131
}
112132

‎pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/NamespaceAuthZTest.java‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,12 @@
3636
import java.util.UUID;
3737
import java.util.concurrent.TimeUnit;
3838
import java.util.concurrent.atomic.AtomicBoolean;
39+
import java.util.function.Consumer;
3940
import lombok.Cleanup;
4041
import lombok.SneakyThrows;
4142
import org.apache.commons.lang3.StringUtils;
4243
import org.apache.commons.lang3.reflect.FieldUtils;
44+
import org.apache.pulsar.broker.BrokerTestUtil;
4345
import org.apache.pulsar.broker.authorization.AuthorizationService;
4446
import org.apache.pulsar.broker.service.Topic;
4547
import org.apache.pulsar.client.admin.PulsarAdmin;
@@ -96,6 +98,9 @@ public class NamespaceAuthZTest extends MockedPulsarStandalone {
9698
private static final String TENANT_ADMIN_TOKEN = Jwts.builder()
9799
.claim("sub", TENANT_ADMIN_SUBJECT).signWith(SECRET_KEY).compact();
98100

101+
private volatile Consumer<InvocationOnMock> allowNamespacePolicyOperationAsyncHandler;
102+
private volatile Consumer<InvocationOnMock> allowNamespaceOperationAsyncHandler;
103+
99104
@SneakyThrows
100105
@BeforeClass
101106
public void setup() {
@@ -118,9 +123,28 @@ public void setup() {
118123
.authentication(new AuthenticationToken(TENANT_ADMIN_TOKEN))
119124
.build();
120125
this.pulsarClient = super.getPulsarService().getClient();
121-
this.authorizationService = Mockito.spy(getPulsarService().getBrokerService().getAuthorizationService());
126+
this.authorizationService = BrokerTestUtil.spyWithoutRecordingInvocations(
127+
getPulsarService().getBrokerService().getAuthorizationService());
122128
FieldUtils.writeField(getPulsarService().getBrokerService(), "authorizationService",
123129
authorizationService, true);
130+
Mockito.doAnswer(invocationOnMock -> {
131+
Consumer<InvocationOnMock> localAllowNamespacePolicyOperationAsyncHandler =
132+
allowNamespacePolicyOperationAsyncHandler;
133+
if (localAllowNamespacePolicyOperationAsyncHandler != null) {
134+
localAllowNamespacePolicyOperationAsyncHandler.accept(invocationOnMock);
135+
}
136+
return invocationOnMock.callRealMethod();
137+
}).when(authorizationService).allowNamespacePolicyOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
138+
Mockito.any(), Mockito.any());
139+
Mockito.doAnswer(invocationOnMock -> {
140+
Consumer<InvocationOnMock> localAllowNamespaceOperationAsyncHandler =
141+
allowNamespaceOperationAsyncHandler;
142+
if (localAllowNamespaceOperationAsyncHandler != null) {
143+
localAllowNamespaceOperationAsyncHandler.accept(invocationOnMock);
144+
}
145+
return invocationOnMock.callRealMethod();
146+
}).when(authorizationService).allowNamespaceOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
147+
Mockito.any());
124148
}
125149

126150

@@ -144,33 +168,31 @@ public void cleanup() {
144168
public void after() throws Exception {
145169
deleteNamespaceWithRetry("public/default", true, superUserAdmin);
146170
superUserAdmin.namespaces().createNamespace("public/default");
171+
allowNamespacePolicyOperationAsyncHandler = null;
172+
allowNamespaceOperationAsyncHandler = null;
147173
}
148174

149175
private AtomicBoolean setAuthorizationOperationChecker(String role, NamespaceOperation operation) {
150176
AtomicBoolean execFlag = new AtomicBoolean(false);
151-
Mockito.doAnswer(invocationOnMock -> {
177+
allowNamespaceOperationAsyncHandler = invocationOnMock -> {
152178
String role_ = invocationOnMock.getArgument(2);
153179
if (role.equals(role_)) {
154180
NamespaceOperation operation_ = invocationOnMock.getArgument(1);
155181
Assert.assertEquals(operation_, operation);
156182
}
157183
execFlag.set(true);
158-
return invocationOnMock.callRealMethod();
159-
}).when(authorizationService).allowNamespaceOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
160-
Mockito.any());
161-
return execFlag;
184+
};
185+
return execFlag;
162186
}
163187

164188
private void clearAuthorizationOperationChecker() {
165-
Mockito.doAnswer(InvocationOnMock::callRealMethod).when(authorizationService)
166-
.allowNamespaceOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
167-
Mockito.any());
189+
allowNamespaceOperationAsyncHandler = null;
168190
}
169191

170192
private AtomicBoolean setAuthorizationPolicyOperationChecker(String role, Object policyName, Object operation) {
171193
AtomicBoolean execFlag = new AtomicBoolean(false);
172194
if (operation instanceof PolicyOperation) {
173-
Mockito.doAnswer(invocationOnMock -> {
195+
allowNamespacePolicyOperationAsyncHandler = invocationOnMock -> {
174196
String role_ = invocationOnMock.getArgument(3);
175197
if (role.equals(role_)) {
176198
PolicyName policyName_ = invocationOnMock.getArgument(1);
@@ -179,9 +201,7 @@ private AtomicBoolean setAuthorizationPolicyOperationChecker(String role, Object
179201
assertEquals(policyName_, policyName);
180202
}
181203
execFlag.set(true);
182-
return invocationOnMock.callRealMethod();
183-
}).when(authorizationService).allowNamespacePolicyOperationAsync(Mockito.any(), Mockito.any(), Mockito.any(),
184-
Mockito.any(), Mockito.any());
204+
};
185205
} else {
186206
throw new IllegalArgumentException("");
187207
}

‎pulsar-broker/src/test/java/org/apache/pulsar/broker/admin/TransactionAndSchemaAuthZTest.java‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import java.util.concurrent.atomic.AtomicBoolean;
2626
import lombok.Cleanup;
2727
import lombok.SneakyThrows;
28-
import org.apache.commons.lang3.reflect.FieldUtils;
2928
import org.apache.pulsar.client.admin.PulsarAdmin;
3029
import org.apache.pulsar.client.admin.PulsarAdminException;
3130
import org.apache.pulsar.client.api.Consumer;
@@ -43,12 +42,9 @@
4342
import org.apache.pulsar.common.schema.SchemaInfo;
4443
import org.apache.pulsar.common.schema.SchemaType;
4544
import org.apache.pulsar.metadata.api.MetadataStoreException;
46-
import org.mockito.Mockito;
4745
import org.testng.Assert;
4846
import org.testng.annotations.AfterClass;
49-
import org.testng.annotations.AfterMethod;
5047
import org.testng.annotations.BeforeClass;
51-
import org.testng.annotations.BeforeMethod;
5248
import org.testng.annotations.DataProvider;
5349
import org.testng.annotations.Test;
5450

@@ -85,20 +81,6 @@ public void cleanup() {
8581
close();
8682
}
8783

88-
@BeforeMethod
89-
public void before() throws IllegalAccessException {
90-
orignalAuthorizationService = getPulsarService().getBrokerService().getAuthorizationService();
91-
authorizationService = Mockito.spy(orignalAuthorizationService);
92-
FieldUtils.writeField(getPulsarService().getBrokerService(), "authorizationService",
93-
authorizationService, true);
94-
}
95-
96-
@AfterMethod
97-
public void after() throws IllegalAccessException {
98-
FieldUtils.writeField(getPulsarService().getBrokerService(), "authorizationService",
99-
orignalAuthorizationService, true);
100-
}
101-
10284
protected void createTransactionCoordinatorAssign(int numPartitionsOfTC) throws MetadataStoreException {
10385
getPulsarService().getPulsarResources()
10486
.getNamespaceResources()

0 commit comments

Comments
 (0)