Fix test setup to properly free context (elastic#130787) (elastic#130835)

Author: benchaplin

muted-tests.yml

@@ -432,12 +432,6 @@ tests:
   - class: org.elasticsearch.action.RejectionActionIT
     method: testSimulatedSearchRejectionLoad
     issue: https://github.com/elastic/elasticsearch/issues/125901
-  - class: org.elasticsearch.search.SearchServiceSingleNodeTests
-    method: testBeforeShardLockDuringShardCreate
-    issue: https://github.com/elastic/elasticsearch/issues/126812
-  - class: org.elasticsearch.search.SearchServiceSingleNodeTests
-    method: testLookUpSearchContext
-    issue: https://github.com/elastic/elasticsearch/issues/126813
   - class: org.elasticsearch.test.index.IndexVersionUtilsTests
     method: testIndexCompatibleVersionMatches
     issue: https://github.com/elastic/elasticsearch/issues/120760

server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java

@@ -44,6 +44,7 @@
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.routing.RecoverySource;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.TestShardRouting;
@@ -1259,8 +1260,22 @@ public void testBeforeShardLockDuringShardCreate() {
             TestShardRouting.newShardRouting(
                 new ShardId(indexService.index(), 0),
                 randomAlphaOfLength(5),
-                randomBoolean(),
-                ShardRoutingState.INITIALIZING
+                true,
+                ShardRoutingState.INITIALIZING,
+                RecoverySource.EmptyStoreRecoverySource.INSTANCE
+            ),
+            indexService.getIndexSettings().getSettings()
+        );
+        assertEquals(1, service.getActiveContexts());
+
+        boolean primary = randomBoolean();
+        service.beforeIndexShardCreated(
+            TestShardRouting.newShardRouting(
+                new ShardId(indexService.index(), 0),
+                randomAlphaOfLength(5),
+                primary,
+                ShardRoutingState.INITIALIZING,
+                primary ? RecoverySource.ExistingStoreRecoverySource.INSTANCE : RecoverySource.PeerRecoverySource.INSTANCE
             ),
             indexService.getIndexSettings().getSettings()
         );

test/framework/src/main/java/org/elasticsearch/cluster/routing/TestShardRouting.java

@@ -127,6 +127,16 @@ public static ShardRouting newShardRouting(String index, int shardId, String cur
         return newShardRouting(new ShardId(index, IndexMetadata.INDEX_UUID_NA_VALUE, shardId), currentNodeId, primary, state);
     }
 
+    public static ShardRouting newShardRouting(
+        ShardId shardId,
+        String currentNodeId,
+        boolean primary,
+        ShardRoutingState state,
+        RecoverySource recoverySource
+    ) {
+        return newShardRouting(shardId, currentNodeId, primary, state, recoverySource, ShardRouting.Role.DEFAULT);
+    }
+
     public static ShardRouting newShardRouting(ShardId shardId, String currentNodeId, boolean primary, ShardRoutingState state) {
         assertNotEquals(ShardRoutingState.RELOCATING, state);
         return new ShardRouting(
@@ -144,6 +154,30 @@ public static ShardRouting newShardRouting(ShardId shardId, String currentNodeId
         );
     }
 
+    public static ShardRouting newShardRouting(
+        ShardId shardId,
+        String currentNodeId,
+        boolean primary,
+        ShardRoutingState state,
+        RecoverySource recoverySource,
+        ShardRouting.Role role
+    ) {
+        assertNotEquals(ShardRoutingState.RELOCATING, state);
+        return new ShardRouting(
+            shardId,
+            currentNodeId,
+            null,
+            primary,
+            state,
+            recoverySource,
+            buildUnassignedInfo(state),
+            buildRelocationFailureInfo(state),
+            buildAllocationId(state),
+            ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE,
+            role
+        );
+    }
+
     public static ShardRouting newShardRouting(
         String index,
         int shardId,