[ML] Fix upgrade test assertion on trained model metadata (elastic#95748

)

The failure was in an assertion that the trained model metadata
has been renamed trained_model_allocation -> 
trained_model_assignment but this can only happen if there is 
a cluster change event after all nodes have been upgraded.
Relaxes the assertion to account for this case.

x-pack/qa/rolling-upgrade/src/test/java/org/elasticsearch/upgrades/MLModelDeploymentsUpgradeIT.java

@@ -26,9 +26,9 @@
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.client.WarningsHandler.PERMISSIVE;
-import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.oneOf;
 
 public class MLModelDeploymentsUpgradeIT extends AbstractUpgradeTestCase {
 
@@ -73,6 +73,7 @@ public void setUpLogging() throws IOException {
             {
               "persistent": {
                 "logger.org.elasticsearch.xpack.ml.inference": "TRACE",
+                "logger.org.elasticsearch.xpack.ml.inference.assignments": "DEBUG",
                 "logger.org.elasticsearch.xpack.ml.process": "DEBUG",
                 "logger.org.elasticsearch.xpack.ml.action": "TRACE"
               }
@@ -133,7 +134,6 @@ public void testTrainedModelDeployment() throws Exception {
         }
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/95501")
     public void testTrainedModelDeploymentStopOnMixedCluster() throws Exception {
         assumeTrue("NLP model deployments added in 8.0", UPGRADE_FROM_VERSION.onOrAfter(Version.V_8_0_0));
 
@@ -157,7 +157,6 @@ public void testTrainedModelDeploymentStopOnMixedCluster() throws Exception {
                     request.addParameter("timeout", "70s");
                 }));
                 assertThatTrainedModelAssignmentMetadataIsEmpty(modelId);
-
             }
             default -> throw new UnsupportedOperationException("Unknown cluster type [" + CLUSTER_TYPE + "]");
         }
@@ -263,11 +262,18 @@ private void assertThatTrainedModelAssignmentMetadataIsEmpty(String modelId) thr
             "_cluster/state?filter_path=metadata.trained_model_assignment." + modelId
         );
         Response getTrainedModelAssignmentMetadataResponse = client().performRequest(getTrainedModelAssignmentMetadataRequest);
-        assertThat(EntityUtils.toString(getTrainedModelAssignmentMetadataResponse.getEntity()), containsString("{}"));
-
+        String responseBody = EntityUtils.toString(getTrainedModelAssignmentMetadataResponse.getEntity());
+        assertThat(responseBody, oneOf("{}", "{\"metadata\":{\"trained_model_assignment\":{}}}"));
+
+        // trained_model_allocation was renamed to trained_model_assignment
+        // in v8.3. The renaming happens automatically and the old
+        // metadata should be removed once all nodes are upgraded.
+        // However, sometimes there aren't enough cluster state change
+        // events in the upgraded cluster test for this to happen
         getTrainedModelAssignmentMetadataRequest = new Request("GET", "_cluster/state?filter_path=metadata.trained_model_allocation");
         getTrainedModelAssignmentMetadataResponse = client().performRequest(getTrainedModelAssignmentMetadataRequest);
-        assertThat(EntityUtils.toString(getTrainedModelAssignmentMetadataResponse.getEntity()), equalTo("{}"));
+        responseBody = EntityUtils.toString(getTrainedModelAssignmentMetadataResponse.getEntity());
+        assertThat(responseBody, oneOf("{}", "{\"metadata\":{\"trained_model_allocation\":{}}}"));
     }
 
     private Response getTrainedModelStats(String modelId) throws IOException {