From 22b236f5d72793a628ea7473c4bca3776a3856cb Mon Sep 17 00:00:00 2001
From: Chen Dai <daichen@amazon.com>
Date: Tue, 2 Jul 2024 12:06:01 -0700
Subject: [PATCH] Add more UT

Signed-off-by: Chen Dai <daichen@amazon.com>
---
 .../ApplyFlintSparkCoveringIndexSuite.scala   | 55 ++++++++++++-------
 1 file changed, 35 insertions(+), 20 deletions(-)

diff --git a/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/covering/ApplyFlintSparkCoveringIndexSuite.scala b/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/covering/ApplyFlintSparkCoveringIndexSuite.scala
index fe1e172d2..af19b3801 100644
--- a/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/covering/ApplyFlintSparkCoveringIndexSuite.scala
+++ b/flint-spark-integration/src/test/scala/org/opensearch/flint/spark/covering/ApplyFlintSparkCoveringIndexSuite.scala
@@ -64,33 +64,48 @@ class ApplyFlintSparkCoveringIndexSuite extends FlintSuite with Matchers {
   }
 
   Seq(
-    /*("age = 30", "age = 20", false),
+    ("age = 30", "age = 20", false),
     ("age = 30", "age < 20", false),
     ("age = 30", "age > 50", false),
     ("age > 30 AND age < 60", "age > 20 AND age < 50", false),
+    ("age > 30", "age >= 30", false),
+    ("age <= 30", "age <= 20", false),
+    ("age < 50", "age = 49", false),
+    ("age <= 50", "age = 50", false),
+    ("age > 30 AND age < 60", "age > 40 AND age < 50", false),
+    (null, "age > 30", false), // no query filter
     ("age = 30", "age = 30", true),
     ("age = 30", "age <= 30", true),
     ("age = 30", "age >= 30", true),
-    ("age = 30", "age > 20 AND age < 50", true),*/
-    ("age > 30 AND age < 40", "age > 20 AND age < 50", true)).foreach {
-    case (queryFilter, indexFilter, expectedResult) =>
-      test(
-        s"apply partial covering index with [$indexFilter] to query filter [$queryFilter]: $expectedResult") {
-        val assertion = assertFlintQueryRewriter
-          .withQuery(s"SELECT name FROM $testTable WHERE $queryFilter")
-          .withIndex(
-            new FlintSparkCoveringIndex(
-              indexName = "partial",
-              tableName = testTable,
-              indexedColumns = Map("name" -> "string", "age" -> "int"),
-              filterCondition = Some(indexFilter)))
-
-        if (expectedResult) {
-          assertion.assertIndexUsed(getFlintIndexName("partial", testTable))
-        } else {
-          assertion.assertIndexNotUsed(testTable)
-        }
+    ("age = 30", "age > 20 AND age < 50", true),
+    ("age > 30 AND age < 40", "age > 20 AND age < 50", true),
+    ("age >= 30", "age > 29", true),
+    ("age <= 30", "age < 31", true),
+    ("age > 30", null, true) // no index filter
+  ).foreach { case (queryFilter, indexFilter, expectedResult) =>
+    test(
+      s"apply partial covering index with [$indexFilter] to query filter [$queryFilter]: $expectedResult") {
+      val query = if (queryFilter == null) {
+        s"SELECT name FROM $testTable"
+      } else {
+        s"SELECT name FROM $testTable WHERE $queryFilter"
       }
+
+      val assertion = assertFlintQueryRewriter
+        .withQuery(query)
+        .withIndex(
+          new FlintSparkCoveringIndex(
+            indexName = "partial",
+            tableName = testTable,
+            indexedColumns = Map("name" -> "string", "age" -> "int"),
+            filterCondition = Option(indexFilter)))
+
+      if (expectedResult) {
+        assertion.assertIndexUsed(getFlintIndexName("partial", testTable))
+      } else {
+        assertion.assertIndexNotUsed(testTable)
+      }
+    }
   }
 
   test("should not apply if covering index is logically deleted") {