Fixup other jmock related test failures

nbauernfeind · Oct 22, 2024 · fbb377f · fbb377f
1 parent bcd681e
commit fbb377f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/...ain/java/io/deephaven/engine/table/impl/sources/regioned/RegionedColumnSourceManager.java b/...ain/java/io/deephaven/engine/table/impl/sources/regioned/RegionedColumnSourceManager.java
@@ -426,11 +426,15 @@ public final synchronized boolean isEmpty() {
     protected void destroy() {
         super.destroy();
         for (final EmptyTableLocationEntry entry : emptyTableLocations.values()) {
-            entry.subscriptionBuffer.reset();
+            if (entry.subscriptionBuffer != null) {
+                entry.subscriptionBuffer.reset();
+            }
         }
         emptyTableLocations.clear();
         for (final IncludedTableLocationEntry entry : includedTableLocations.values()) {
-            entry.subscriptionBuffer.reset();
+            if (entry.subscriptionBuffer != null) {
+                entry.subscriptionBuffer.reset();
+            }
         }
         includedTableLocations.clear();
     }

diff --git a/engine/table/src/test/java/io/deephaven/engine/table/impl/TestPartitionAwareSourceTable.java b/engine/table/src/test/java/io/deephaven/engine/table/impl/TestPartitionAwareSourceTable.java
@@ -165,6 +165,7 @@ public void setUp() throws Exception {
     @Override
     public void tearDown() throws Exception {
         try {
+            allowLivenessRelease();
             super.tearDown();
         } finally {
             if (coalesced != null) {
@@ -174,6 +175,22 @@ public void tearDown() throws Exception {
         }
     }
 
+    private void allowLivenessRelease() {
+        checking(new Expectations() {
+            {
+                allowing(locationProvider).supportsSubscriptions();
+                allowing(locationProvider).unsubscribe(with(any(TableLocationProvider.Listener.class)));
+                will(returnValue(true));
+                for (int li = 0; li < tableLocations.length; ++li) {
+                    final TableLocation tableLocation = tableLocations[li];
+                    allowing(tableLocation).supportsSubscriptions();
+                    will(returnValue(true));
+                    allowing(tableLocation).unsubscribe(with(any(TableLocation.Listener.class)));
+                }
+            }
+        });
+    }
+
     private Map<String, ? extends ColumnSource<?>> getIncludedColumnsMap(final int... indices) {
         return IntStream.of(indices)
                 .mapToObj(ci -> new Pair<>(TABLE_DEFINITION.getColumns().get(ci).getName(), columnSources[ci]))
@@ -396,6 +413,7 @@ public Object invoke(Invocation invocation) {
         errorNotification.reset();
         final ControlledUpdateGraph updateGraph = ExecutionContext.getContext().getUpdateGraph().cast();
         updateGraph.runWithinUnitTestCycle(() -> {
+            allowLivenessRelease();
             SUT.refresh();
             updateGraph.markSourcesRefreshedForUnitTests();
         }, false);