From 93d8a8a8ebf0a97d9a48b273de058ac228981c91 Mon Sep 17 00:00:00 2001
From: Ian Hoang <51065478+IanHoang@users.noreply.github.com>
Date: Mon, 12 Jun 2023 15:11:34 -0500
Subject: [PATCH] Update Codec Enum to support ZSTD and ZSTDNODICT (#327)

Signed-off-by: Ian Hoang <hoangia@amazon.com>
Co-authored-by: Ian Hoang <hoangia@amazon.com>
---
 osbenchmark/workload/workload.py | 14 +++++---
 tests/workload/params_test.py    | 59 +++++++++++++++++++++++++++++++-
 2 files changed, 68 insertions(+), 5 deletions(-)

diff --git a/osbenchmark/workload/workload.py b/osbenchmark/workload/workload.py
index 3331d5598..a92d732b9 100644
--- a/osbenchmark/workload/workload.py
+++ b/osbenchmark/workload/workload.py
@@ -717,14 +717,20 @@ def from_hyphenated_string(cls, v):
 class IndexCodec(Enum):
     Default = "default"
     BestCompression = "best_compression"
+    ZSTD = "zstd"
+    ZSTDNODICT = "zstdnodict"
 
     @classmethod
     def is_codec_valid(cls, codec):
-        for valid_codec in cls:
-            if codec == valid_codec.value:
-                return True
+        available_codecs = cls.get_available_codecs()
+        if codec.lower() in available_codecs:
+            return True
 
-        raise ValueError(f"Invalid index.codec value '{codec}'")
+        raise ValueError(f"Invalid index.codec value '{codec}'. Choose from available codecs: {available_codecs}")
+
+    @classmethod
+    def get_available_codecs(cls):
+        return list(map(lambda codec: codec.value, cls))
 
 
 class TaskNameFilter:
diff --git a/tests/workload/params_test.py b/tests/workload/params_test.py
index 7e14ddb91..a39c79d86 100644
--- a/tests/workload/params_test.py
+++ b/tests/workload/params_test.py
@@ -1632,6 +1632,62 @@ def test_create_index_with_best_compression_codec(self):
         self.assertEqual({}, p["request-params"])
         self.assertEqual("best_compression", body["settings"]["index.codec"])
 
+    def test_create_index_with_zstd_codec(self):
+        source = params.CreateIndexParamSource(workload.Workload(name="unit-test"), params={
+            "index": "test",
+            "body": {
+                "settings": {
+                    "index.number_of_replicas": 0,
+                    "index.codec": "zstd"
+                },
+                "mappings": {
+                    "doc": {
+                        "properties": {
+                            "name": {
+                                "type": "keyword",
+                            }
+                        }
+                    }
+                }
+            }
+        })
+
+        p = source.params()
+        self.assertEqual(1, len(p["indices"]))
+        index, body = p["indices"][0]
+        self.assertEqual("test", index)
+        self.assertTrue(len(body) > 0)
+        self.assertEqual({}, p["request-params"])
+        self.assertEqual("zstd", body["settings"]["index.codec"])
+
+    def test_create_index_with_zstdnodict_codec(self):
+        source = params.CreateIndexParamSource(workload.Workload(name="unit-test"), params={
+            "index": "test",
+            "body": {
+                "settings": {
+                    "index.number_of_replicas": 0,
+                    "index.codec": "zstdnodict"
+                },
+                "mappings": {
+                    "doc": {
+                        "properties": {
+                            "name": {
+                                "type": "keyword",
+                            }
+                        }
+                    }
+                }
+            }
+        })
+
+        p = source.params()
+        self.assertEqual(1, len(p["indices"]))
+        index, body = p["indices"][0]
+        self.assertEqual("test", index)
+        self.assertTrue(len(body) > 0)
+        self.assertEqual({}, p["request-params"])
+        self.assertEqual("zstdnodict", body["settings"]["index.codec"])
+
     def test_create_index_with_invalid_codec(self):
         with self.assertRaises(exceptions.InvalidSyntax) as context:
             params.CreateIndexParamSource(workload.Workload(name="unit-test"), params={
@@ -1654,7 +1710,8 @@ def test_create_index_with_invalid_codec(self):
             })
 
         self.assertEqual(str(context.exception),
-                         "Please set the value properly for the create-index operation. Invalid index.codec value 'invalid_codec'")
+                         "Please set the value properly for the create-index operation. Invalid index.codec value " +
+                         "'invalid_codec'. Choose from available codecs: ['default', 'best_compression', 'zstd', 'zstdnodict']")
 
 class CreateDataStreamParamSourceTests(TestCase):
     def test_create_data_stream(self):