From 66c91c32cfb52c1d87f64b9a9a01421ffbd7ffba Mon Sep 17 00:00:00 2001
From: Thomas Farr <tsfarr@amazon.com>
Date: Tue, 14 May 2024 12:53:34 +1200
Subject: [PATCH] Separate neural query cluster configuration

Signed-off-by: Thomas Farr <tsfarr@amazon.com>
---
 .github/workflows/integration.yml             |  5 ++
 .../Clusters/WritableCluster.cs               |  1 -
 .../Neural/NeuralQueryUsageTests.cs           | 49 ++++++++++---------
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
index 51b4eb0492..73a58cfb99 100644
--- a/.github/workflows/integration.yml
+++ b/.github/workflows/integration.yml
@@ -51,6 +51,11 @@ jobs:
         name: Integration Tests
         working-directory: client
 
+      # Run neural query integration tests separately as they use a significant amount of memory on their own
+      - run: "./build.sh integrate ${{ matrix.version }} neuralquery random:test_only_one --report"
+        name: Neural Query Integration Tests
+        working-directory: client
+
       - name: Upload test report
         if: failure()
         uses: actions/upload-artifact@v3
diff --git a/tests/Tests.Core/ManagedOpenSearch/Clusters/WritableCluster.cs b/tests/Tests.Core/ManagedOpenSearch/Clusters/WritableCluster.cs
index 145cc6b15c..320aab24ea 100644
--- a/tests/Tests.Core/ManagedOpenSearch/Clusters/WritableCluster.cs
+++ b/tests/Tests.Core/ManagedOpenSearch/Clusters/WritableCluster.cs
@@ -49,7 +49,6 @@ private static ClientTestClusterConfiguration CreateConfiguration() =>
 				Knn,
                 MachineLearning,
 				MapperMurmur3,
-                NeuralSearch,
 				Security)
 			{
 				MaxConcurrency = 4,
diff --git a/tests/Tests/QueryDsl/Specialized/Neural/NeuralQueryUsageTests.cs b/tests/Tests/QueryDsl/Specialized/Neural/NeuralQueryUsageTests.cs
index 4a67c3cd08..28745bc74c 100644
--- a/tests/Tests/QueryDsl/Specialized/Neural/NeuralQueryUsageTests.cs
+++ b/tests/Tests/QueryDsl/Specialized/Neural/NeuralQueryUsageTests.cs
@@ -6,13 +6,13 @@
 */
 
 using System;
-using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using FluentAssertions;
 using OpenSearch.Client;
 using OpenSearch.Net;
 using OpenSearch.OpenSearch.Xunit.XunitPlumbing;
+using OpenSearch.Stack.ArtifactsApi.Products;
 using Tests.Core.Extensions;
 using Tests.Core.ManagedOpenSearch.Clusters;
 using Tests.Framework.EndpointTests.TestState;
@@ -20,6 +20,30 @@
 
 namespace Tests.QueryDsl.Specialized.Neural;
 
+public class NeuralQueryCluster : ClientTestClusterBase
+{
+    public NeuralQueryCluster() : base(CreateConfiguration()) { }
+
+    private static ClientTestClusterConfiguration CreateConfiguration()
+    {
+        var config = new ClientTestClusterConfiguration(
+            OpenSearchPlugin.Knn,
+            OpenSearchPlugin.MachineLearning,
+            OpenSearchPlugin.NeuralSearch,
+            OpenSearchPlugin.Security)
+        {
+            MaxConcurrency = 4,
+            ValidatePluginsToInstall = false,
+        };
+
+        config.DefaultNodeSettings.Add("plugins.ml_commons.only_run_on_ml_node", "false");
+        config.DefaultNodeSettings.Add("plugins.ml_commons.native_memory_threshold", "99");
+        config.DefaultNodeSettings.Add("plugins.ml_commons.model_access_control_enabled", "true", ">=2.8.0");
+
+        return config;
+    }
+}
+
 public class NeuralSearchDoc
 {
     [PropertyName("id")] public string Id { get; set; }
@@ -28,16 +52,15 @@ public class NeuralSearchDoc
 }
 
 [SkipVersion("<2.6.0", "Avoid the various early permutations of the ML APIs")]
-[SkipPrereleaseVersions("Prerelease versions of OpenSearch do not include the ML & Neural Search plugins")]
 public class NeuralQueryUsageTests
-    : QueryDslUsageTestsBase<WritableCluster, NeuralSearchDoc>
+    : QueryDslUsageTestsBase<NeuralQueryCluster, NeuralSearchDoc>
 {
     private static readonly string TestName = nameof(NeuralQueryUsageTests).ToLowerInvariant();
 
     private string _modelGroupId;
     private string _modelId = "default-for-unit-tests";
 
-    public NeuralQueryUsageTests(WritableCluster cluster, EndpointUsage usage) : base(cluster, usage) { }
+    public NeuralQueryUsageTests(NeuralQueryCluster cluster, EndpointUsage usage) : base(cluster, usage) { }
 
     protected override IndexName IndexName => TestName;
     protected override string ExpectedIndexString => TestName;
@@ -129,24 +152,6 @@ protected override void IntegrationSetup(IOpenSearchClient client, CallUniqueVal
         var renamedToRegisterDeploy = baseVersion >= new Version("2.7.0");
         var hasModelAccessControl = baseVersion >= new Version("2.8.0");
 
-        var settings = new Dictionary<string, object>
-        {
-            ["plugins.ml_commons.only_run_on_ml_node"] = false,
-            ["plugins.ml_commons.native_memory_threshold"] = 99
-        };
-
-        if (hasModelAccessControl)
-            settings["plugins.ml_commons.model_access_control_enabled"] = true;
-
-        if (settings.Count > 0)
-        {
-            var putSettingsResp = client.Cluster.PutSettings(new ClusterPutSettingsRequest
-            {
-                Transient = settings
-            });
-            putSettingsResp.ShouldBeValid();
-        }
-
         if (hasModelAccessControl)
         {
             var registerModelGroupResp = client.Http.Post<DynamicResponse>(