fix ut

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

ml-algorithms/src/test/java/org/opensearch/ml/engine/httpclient/MLHttpClientFactoryTests.java

@@ -8,6 +8,7 @@
 import static org.junit.Assert.assertNotNull;
 
 import java.time.Duration;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.Rule;
 import org.junit.Test;
@@ -28,70 +29,84 @@ public void test_getSdkAsyncHttpClient_success() {
 
     @Test
     public void test_validateIp_validIp_noException() throws Exception {
-        MLHttpClientFactory.validate("http", "api.openai.com", 80);
+        AtomicBoolean privateIpEnabled = new AtomicBoolean(false);
+        MLHttpClientFactory.validate("http", "api.openai.com", 80, privateIpEnabled);
     }
 
     @Test
     public void test_validateIp_rarePrivateIp_throwException() throws Exception {
+        AtomicBoolean privateIpEnabled = new AtomicBoolean(false);
         try {
-            MLHttpClientFactory.validate("http", "0254.020.00.01", 80);
+            MLHttpClientFactory.validate("http", "0254.020.00.01", 80, privateIpEnabled);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
 
         try {
-            MLHttpClientFactory.validate("http", "172.1048577", 80);
+            MLHttpClientFactory.validate("http", "172.1048577", 80, privateIpEnabled);
         } catch (Exception e) {
             assertNotNull(e);
         }
 
         try {
-            MLHttpClientFactory.validate("http", "2886729729", 80);
+            MLHttpClientFactory.validate("http", "2886729729", 80, privateIpEnabled);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
 
         try {
-            MLHttpClientFactory.validate("http", "192.11010049", 80);
+            MLHttpClientFactory.validate("http", "192.11010049", 80, privateIpEnabled);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
 
         try {
-            MLHttpClientFactory.validate("http", "3232300545", 80);
+            MLHttpClientFactory.validate("http", "3232300545", 80, privateIpEnabled);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
 
         try {
-            MLHttpClientFactory.validate("http", "0:0:0:0:0:ffff:127.0.0.1", 80);
+            MLHttpClientFactory.validate("http", "0:0:0:0:0:ffff:127.0.0.1", 80, privateIpEnabled);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
 
         try {
-            MLHttpClientFactory.validate("http", "153.24.76.232", 80);
+            MLHttpClientFactory.validate("http", "153.24.76.232", 80, privateIpEnabled);
         } catch (IllegalArgumentException e) {
             assertNotNull(e);
         }
     }
 
+    @Test
+    public void test_validateIp_rarePrivateIp_NotThrowException() throws Exception {
+        AtomicBoolean privateIpEnabled = new AtomicBoolean(true);
+        MLHttpClientFactory.validate("http", "0254.020.00.01", 80, privateIpEnabled);
+        MLHttpClientFactory.validate("http", "172.1048577", 80, privateIpEnabled);
+        MLHttpClientFactory.validate("http", "2886729729", 80, privateIpEnabled);
+        MLHttpClientFactory.validate("http", "192.11010049", 80, privateIpEnabled);
+        MLHttpClientFactory.validate("http", "3232300545", 80, privateIpEnabled);
+        MLHttpClientFactory.validate("http", "0:0:0:0:0:ffff:127.0.0.1", 80, privateIpEnabled);
+        MLHttpClientFactory.validate("http", "153.24.76.232", 80, privateIpEnabled);
+    }
+
     @Test
     public void test_validateSchemaAndPort_success() throws Exception {
-        MLHttpClientFactory.validate("http", "api.openai.com", 80);
+        MLHttpClientFactory.validate("http", "api.openai.com", 80, new AtomicBoolean(false));
     }
 
     @Test
     public void test_validateSchemaAndPort_notAllowedSchema_throwException() throws Exception {
         expectedException.expect(IllegalArgumentException.class);
-        MLHttpClientFactory.validate("ftp", "api.openai.com", 80);
+        MLHttpClientFactory.validate("ftp", "api.openai.com", 80, new AtomicBoolean(false));
     }
 
     @Test
     public void test_validateSchemaAndPort_portNotInRange_throwException() throws Exception {
         expectedException.expect(IllegalArgumentException.class);
         expectedException.expectMessage("Port out of range: 65537");
-        MLHttpClientFactory.validate("https", "api.openai.com", 65537);
+        MLHttpClientFactory.validate("https", "api.openai.com", 65537, new AtomicBoolean(false));
     }
 
 }

plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java

@@ -61,6 +61,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 
 import org.junit.Before;
@@ -109,6 +110,7 @@
 import org.opensearch.ml.engine.encryptor.Encryptor;
 import org.opensearch.ml.engine.encryptor.EncryptorImpl;
 import org.opensearch.ml.engine.indices.MLIndicesHandler;
+import org.opensearch.ml.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.stats.ActionName;
 import org.opensearch.ml.stats.MLActionLevelStat;
 import org.opensearch.ml.stats.MLNodeLevelStat;
@@ -179,6 +181,8 @@ public class MLModelManagerTests extends OpenSearchTestCase {
 
     @Mock
     ClusterApplierService clusterApplierService;
+    @Mock
+    MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     @Before
     public void setup() throws URISyntaxException {
@@ -253,6 +257,8 @@ public void setup() throws URISyntaxException {
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
 
+        when(mlFeatureEnabledSetting.isConnectorPrivateIpEnabled()).thenReturn(new AtomicBoolean(false));
+
         modelManager = spy(
             new MLModelManager(
                 clusterService,
@@ -268,7 +274,8 @@ public void setup() throws URISyntaxException {
                 mlTaskManager,
                 modelCacheHelper,
                 mlEngine,
-                nodeHelper
+                nodeHelper,
+                mlFeatureEnabledSetting
             )
         );