Fix test setup to properly free context

Author: benchaplin

muted-tests.yml

@@ -323,12 +323,6 @@ tests:
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=transform/transforms_start_stop/Test start/stop/start continuous transform}
   issue: https://github.com/elastic/elasticsearch/issues/126755
-- 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.xpack.test.rest.XPackRestIT
   method: test {p0=transform/transforms_stats/Test get multiple transform stats where one does not have a task}
   issue: https://github.com/elastic/elasticsearch/issues/126863

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;
@@ -1255,8 +1256,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

@@ -131,6 +131,16 @@ public static ShardRouting newShardRouting(ShardId shardId, String currentNodeId
         return newShardRouting(shardId, currentNodeId, primary, state, ShardRouting.Role.DEFAULT);
     }
 
+    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,
@@ -154,6 +164,30 @@ public static ShardRouting newShardRouting(
         );
     }
 
+    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,