[8.x] [ML] Fix ModelRegistryMetadataTests intermittent failures (#125883) (#126042)

* [ML] Fix ModelRegistryMetadataTests intermittent failures (#125883)

unmute test

(cherry picked from commit 5582794)

# Conflicts:
#	muted-tests.yml
#	x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/registry/ModelRegistryMetadata.java

* Unmute testAlreadyUpgraded

---------

Co-authored-by: Jim Ferenczi <jim.ferenczi@elastic.co>

Author: Mikep86

muted-tests.yml

@@ -468,9 +468,6 @@ tests:
 - class: org.elasticsearch.qa.verify_version_constants.VerifyVersionConstantsIT
   method: testLuceneVersionConstant
   issue: https://github.com/elastic/elasticsearch/issues/125638
-- class: org.elasticsearch.xpack.inference.registry.ModelRegistryMetadataTests
-  method: testUpgrade
-  issue: https://github.com/elastic/elasticsearch/issues/125554
 - class: org.elasticsearch.xpack.application.HuggingFaceServiceUpgradeIT
   method: testElser {upgradedNodes=2}
   issue: https://github.com/elastic/elasticsearch/issues/125550
@@ -492,9 +489,6 @@ tests:
 - class: org.elasticsearch.packaging.test.DebPreservationTests
   method: test40RestartOnUpgrade
   issue: https://github.com/elastic/elasticsearch/issues/125821
-- class: org.elasticsearch.xpack.inference.registry.ModelRegistryMetadataTests
-  method: testAlreadyUpgraded
-  issue: https://github.com/elastic/elasticsearch/issues/125585
 - class: org.elasticsearch.xpack.application.HuggingFaceServiceUpgradeIT
   method: testHFEmbeddings {upgradedNodes=2}
   issue: https://github.com/elastic/elasticsearch/issues/125551

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

@@ -47,7 +47,8 @@
 public class ModelRegistryMetadata implements Metadata.Custom {
     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");
@@ -83,7 +84,7 @@ public class ModelRegistryMetadata implements Metadata.Custom {
 
     public static ModelRegistryMetadata fromState(Metadata metadata) {
         ModelRegistryMetadata resp = metadata.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"));
     }
 }