Fix unit test after replace pair

NationalSecurityAgency · Sep 25, 2024 · ec2ea59 · ec2ea59
1 parent 2fa9af9
commit ec2ea59
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/web-services/query/src/test/java/datawave/webservice/query/cache/QueryCacheBeanTest.java b/web-services/query/src/test/java/datawave/webservice/query/cache/QueryCacheBeanTest.java
@@ -56,9 +56,9 @@ public void testListRunningQueries() {
         // Set expectations
         expect(altCache.iterator()).andReturn((Iterator<RunningQuery>) new HashMap().values().iterator());
         Map<String,AbstractMap.SimpleEntry<QueryLogic<?>,AccumuloClient>> snapshot = new HashMap<>();
+        this.pair = new AbstractMap.SimpleEntry<>(this.logic, null);
         snapshot.put("key", this.pair);
         expect(this.remoteCache.snapshot()).andReturn(snapshot);
-        expect(this.pair.getKey()).andReturn((QueryLogic) this.logic);
 
         // Run the test
         PowerMock.replayAll();

diff --git a/...s/query/src/test/java/datawave/webservice/query/runner/ExtendedQueryExecutorBeanTest.java b/...s/query/src/test/java/datawave/webservice/query/runner/ExtendedQueryExecutorBeanTest.java
@@ -244,6 +244,8 @@ public void setupBefore() throws Exception {
         queryExpirationConf.setShortCircuitCheckTime(45);
         queryExpirationConf.setShortCircuitTimeout(58);
         queryExpirationConf.setIdleTimeout(60);
+
+        this.tuple = new AbstractMap.SimpleEntry<>(this.queryLogic1, this.client);
     }
 
     @SuppressWarnings({"rawtypes", "unchecked"})
@@ -254,9 +256,7 @@ public void testAdminCancel_HappyPath() throws Exception {
         // Set expectations of the create logic
         expect(this.connectionRequestBean.adminCancelConnectionRequest(queryId.toString())).andReturn(false);
         expect(this.qlCache.poll(queryId.toString())).andReturn(this.tuple);
-        expect(this.tuple.getKey()).andReturn((QueryLogic) this.queryLogic1);
         this.queryLogic1.close();
-        expect(this.tuple.getValue()).andReturn(this.client);
         this.connectionFactory.returnClient(this.client);
 
         // Run the test
@@ -418,9 +418,7 @@ public void testAdminClose_NonNullTupleReturned() throws Exception {
         // Set expectations
         expect(this.connectionRequestBean.adminCancelConnectionRequest(queryId.toString())).andReturn(false);
         expect(this.qlCache.poll(queryId.toString())).andReturn(this.tuple);
-        expect(this.tuple.getKey()).andReturn((QueryLogic) this.queryLogic1);
         this.queryLogic1.close();
-        expect(this.tuple.getValue()).andReturn(this.client);
         this.connectionFactory.returnClient(this.client);
 
         // Run the test
@@ -485,11 +483,10 @@ public void testCancel_HappyPath() throws Exception {
         expect(this.principal.getDNs()).andReturn(new String[] {userName});
         expect(this.principal.getShortName()).andReturn(userSid);
         expect(this.principal.getProxyServers()).andReturn(new ArrayList<>(0)).anyTimes();
+
         expect(this.qlCache.pollIfOwnedBy(queryId.toString(), userSid)).andReturn(this.tuple);
         this.closedCache.remove(queryId.toString());
-        expect(this.tuple.getKey()).andReturn((QueryLogic) this.queryLogic1);
         this.queryLogic1.close();
-        expect(this.tuple.getValue()).andReturn(this.client);
         this.connectionFactory.returnClient(this.client);
 
         // Run the test
@@ -663,10 +660,8 @@ public void testClose_UncheckedException() throws Exception {
         expect(this.principal.getProxyServers()).andReturn(new ArrayList<>(0)).anyTimes();
         expect(this.qlCache.pollIfOwnedBy(queryId.toString(), userSid)).andReturn(this.tuple);
         expect(this.principal.getUserDN()).andReturn(SubjectIssuerDNPair.of(userName));
-        expect(this.tuple.getKey()).andReturn((QueryLogic) this.queryLogic1);
         this.queryLogic1.close();
         PowerMock.expectLastCall().andThrow(ILLEGAL_STATE_EXCEPTION);
-        expect(this.tuple.getValue()).andThrow(ILLEGAL_STATE_EXCEPTION);
 
         // Run the test
         PowerMock.replayAll();