Test coverage

Signed-off-by: Andy Kwok <[email protected]>
opensearch-project · Jan 11, 2025 · da7af99 · da7af99
1 parent 302ed85
commit da7af99
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 21 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/expression/ip/IPFunctions.java b/core/src/main/java/org/opensearch/sql/expression/ip/IPFunctions.java
@@ -28,7 +28,6 @@ public class IPFunctions {
 
   public void register(BuiltinFunctionRepository repository) {
     repository.register(cidrmatch());
-    repository.register(geoIp());
   }
 
   private DefaultFunctionResolver cidrmatch() {
@@ -37,18 +36,6 @@ private DefaultFunctionResolver cidrmatch() {
         impl(nullMissingHandling(IPFunctions::exprCidrMatch), BOOLEAN, IP, STRING));
   }
 
-  private DefaultFunctionResolver geoIp() {
-    return define(
-        BuiltinFunctionName.GEOIP.getName(),
-        impl(nullMissingHandling(IPFunctions::exprGeoIp), BOOLEAN, STRING, STRING),
-        impl(
-            nullMissingHandling(IPFunctions::exprGeoIpWithOption),
-            BOOLEAN,
-            STRING,
-            STRING,
-            STRING));
-  }
-
   /**
    * Returns whether the given IP address is within the specified inclusive CIDR IP address range.
    * Supports both IPv4 and IPv6 addresses.
@@ -71,12 +58,4 @@ private ExprValue exprCidrMatch(ExprValue addressExprValue, ExprValue rangeExprV
         : ExprValueUtils.LITERAL_TRUE;
   }
 
-  private ExprValue exprGeoIpWithOption(
-      ExprValue dataSource, ExprValue ipAddress, ExprValue option) {
-    return ExprValueUtils.booleanValue(true);
-  }
-
-  private ExprValue exprGeoIp(ExprValue dataSource, ExprValue ipAddress) {
-    return ExprValueUtils.booleanValue(true);
-  }
 }
diff --git a/core/src/test/java/org/opensearch/sql/expression/ip/GeoIPFunctionTest.java b/core/src/test/java/org/opensearch/sql/expression/ip/GeoIPFunctionTest.java
@@ -5,8 +5,10 @@
 
 package org.opensearch.sql.expression.ip;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 
 import org.junit.jupiter.api.Test;
@@ -16,6 +18,7 @@
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.env.Environment;
 
 @ExtendWith(MockitoExtension.class)
@@ -34,4 +37,10 @@ public void geoIpDefaultImplementation() {
                     .valueOf(env));
     assertTrue(exception.getMessage().matches(".*no default implementation available"));
   }
+
+  @Test
+  public void testGeoipFnctionSignature() {
+    var geoip = DSL.geoip(DSL.literal("HARDCODED_DATASOURCE_NAME"), DSL.ref("ip_address", STRING));
+    assertEquals(BOOLEAN, geoip.type());
+  }
 }
diff --git a/...va/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java b/...va/org/opensearch/sql/opensearch/executor/protector/OpenSearchExecutionProtectorTest.java
@@ -47,6 +47,7 @@
 import org.opensearch.sql.ast.tree.Trendline;
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.data.model.ExprBooleanValue;
+import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.NamedExpression;
@@ -58,10 +59,12 @@
 import org.opensearch.sql.expression.window.ranking.RankFunction;
 import org.opensearch.sql.monitor.ResourceMonitor;
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
+import org.opensearch.sql.opensearch.data.type.OpenSearchTextType;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
 import org.opensearch.sql.opensearch.planner.physical.ADOperator;
 import org.opensearch.sql.opensearch.planner.physical.MLCommonsOperator;
 import org.opensearch.sql.opensearch.planner.physical.MLOperator;
+import org.opensearch.sql.opensearch.planner.physical.OpenSearchEvalOperator;
 import org.opensearch.sql.opensearch.request.OpenSearchRequest;
 import org.opensearch.sql.opensearch.request.OpenSearchRequestBuilder;
 import org.opensearch.sql.opensearch.setting.OpenSearchSettings;
@@ -339,6 +342,25 @@ public void test_visitTrendline() {
         executionProtector.visitTrendline(trendlineOperator, null));
   }
 
+  @Test
+  void test_visitOpenSearchEval() {
+    NodeClient nodeClient = mock(NodeClient.class);
+    OpenSearchEvalOperator evalOperator =
+//    ADOperator adOperator =
+            new OpenSearchEvalOperator(
+                    values(emptyList()),
+                    List.of(
+                            ImmutablePair.of(
+                                    new ReferenceExpression("ageInAbs", OpenSearchTextType.of()),
+                                    DSL.abs(DSL.abs(new ReferenceExpression("age", ExprCoreType.LONG))))
+                    ),
+                    nodeClient);
+
+    assertEquals(
+            executionProtector.doProtect(evalOperator), executionProtector.visitEval(evalOperator, null));
+  }
+
+
   PhysicalPlan resourceMonitor(PhysicalPlan input) {
     return new ResourceMonitorPlan(input, resourceMonitor);
   }