[ML] Fix ModelRegistryMetadataTests intermittent failures (#125883)

unmute test

Author: jimczi

muted-tests.yml

@@ -389,9 +389,6 @@ tests:
 - class: org.elasticsearch.xpack.ilm.TimeSeriesDataStreamsIT
   method: testSearchableSnapshotAction
   issue: https://github.com/elastic/elasticsearch/issues/125867
-- class: org.elasticsearch.xpack.inference.registry.ModelRegistryMetadataTests
-  method: testUpgrade
-  issue: https://github.com/elastic/elasticsearch/issues/125554
 - class: org.elasticsearch.xpack.downsample.DataStreamLifecycleDownsampleDisruptionIT
   method: testDataStreamLifecycleDownsampleRollingRestart
   issue: https://github.com/elastic/elasticsearch/issues/123769

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/registry/ModelRegistryMetadata.java

@@ -51,7 +51,8 @@
 public class ModelRegistryMetadata implements Metadata.ProjectCustom {
     public static final String TYPE = "model_registry";
 
-    public static final ModelRegistryMetadata EMPTY = new ModelRegistryMetadata(ImmutableOpenMap.of(), Set.of());
+    public static final ModelRegistryMetadata EMPTY_NOT_UPGRADED = new ModelRegistryMetadata(ImmutableOpenMap.of(), Set.of());
+    public static final ModelRegistryMetadata EMPTY_UPGRADED = new ModelRegistryMetadata(ImmutableOpenMap.of());
 
     private static final ParseField UPGRADED_FIELD = new ParseField("upgraded");
     private static final ParseField MODELS_FIELD = new ParseField("models");
@@ -87,7 +88,7 @@ public class ModelRegistryMetadata implements Metadata.ProjectCustom {
 
     public static ModelRegistryMetadata fromState(ProjectMetadata projectMetadata) {
         ModelRegistryMetadata resp = projectMetadata.custom(TYPE);
-        return resp != null ? resp : EMPTY;
+        return resp != null ? resp : EMPTY_NOT_UPGRADED;
     }
 
     public ModelRegistryMetadata withAddedModel(String inferenceEntityId, MinimalServiceSettings settings) {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/registry/ModelRegistryMetadataTests.java

@@ -22,16 +22,17 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 
 public class ModelRegistryMetadataTests extends AbstractChunkedSerializingTestCase<ModelRegistryMetadata> {
     public static ModelRegistryMetadata randomInstance() {
-        return randomInstance(randomBoolean());
+        return randomInstance(randomBoolean(), true);
     }
 
-    public static ModelRegistryMetadata randomInstance(boolean isUpgraded) {
-        if (rarely() && isUpgraded == false) {
-            return ModelRegistryMetadata.EMPTY;
+    public static ModelRegistryMetadata randomInstance(boolean isUpgraded, boolean acceptsEmpty) {
+        if (rarely() && acceptsEmpty) {
+            return isUpgraded ? ModelRegistryMetadata.EMPTY_UPGRADED : ModelRegistryMetadata.EMPTY_NOT_UPGRADED;
         }
         int size = randomIntBetween(1, 5);
 
@@ -82,7 +83,7 @@ protected Writeable.Reader<ModelRegistryMetadata> instanceReader() {
     }
 
     public void testUpgrade() {
-        var metadata = randomInstance(false);
+        var metadata = randomInstance(false, false);
         var metadataWithTombstones = metadata.withRemovedModel(Set.of(randomFrom(metadata.getModelMap().keySet())));
 
         var indexMetadata = metadata.withAddedModel(randomAlphanumericOfLength(10), MinimalServiceSettingsTests.randomInstance());
@@ -99,9 +100,10 @@ public void testUpgrade() {
     }
 
     public void testAlreadyUpgraded() {
-        var metadata = randomInstance(true);
-        var indexMetadata = randomInstance(true);
+        var metadata = randomInstance(true, true);
+        var indexMetadata = randomInstance(true, true);
 
         var exc = expectThrows(IllegalArgumentException.class, () -> metadata.withUpgradedModels(indexMetadata.getModelMap()));
+        assertThat(exc.getMessage(), containsString("upgraded"));
     }
 }