From f56ae080bce3cffe313831071c84c6f22a305316 Mon Sep 17 00:00:00 2001
From: Govind Kamat <govkamat@amazon.com>
Date: Wed, 6 Nov 2024 16:52:51 +0000
Subject: [PATCH] Fix workload parameter processing when a key is mapped to a
 JSON file.

Signed-off-by: Govind Kamat <govkamat@amazon.com>
---
 osbenchmark/utils/opts.py      |  2 +-
 osbenchmark/workload/loader.py |  6 +++++-
 tests/utils/opts_test.py       | 21 +++++++++++++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/osbenchmark/utils/opts.py b/osbenchmark/utils/opts.py
index fcc029b60..784ad7549 100644
--- a/osbenchmark/utils/opts.py
+++ b/osbenchmark/utils/opts.py
@@ -84,7 +84,7 @@ def convert(v):
 
 
 def to_dict(arg, default_parser=kv_to_map):
-    if io.has_extension(arg, ".json"):
+    if io.has_extension(arg, ".json") and ',' not in arg and ':' not in arg:
         with open(io.normalize_path(arg), mode="rt", encoding="utf-8") as f:
             return json.load(f)
     elif arg.startswith("{"):
diff --git a/osbenchmark/workload/loader.py b/osbenchmark/workload/loader.py
index a97f85fe3..a57c3d901 100644
--- a/osbenchmark/workload/loader.py
+++ b/osbenchmark/workload/loader.py
@@ -1278,7 +1278,11 @@ def read(self, workload_name, workload_spec_file, mapping_dir):
                         ve.instance, indent=4, sort_keys=True),
                         ve.absolute_path, ve.absolute_schema_path))
 
-        current_workload = self.read_workload(workload_name, workload_spec, mapping_dir)
+        try:
+            current_workload = self.read_workload(workload_name, workload_spec, mapping_dir)
+        except Exception as e:
+            console.error(e)
+            raise
 
         unused_user_defined_workload_params = self.complete_workload_params.unused_user_defined_workload_params()
         if len(unused_user_defined_workload_params) > 0:
diff --git a/tests/utils/opts_test.py b/tests/utils/opts_test.py
index 617006a06..4f4febd79 100644
--- a/tests/utils/opts_test.py
+++ b/tests/utils/opts_test.py
@@ -25,6 +25,7 @@
 import os
 from unittest import TestCase
 
+import unittest.mock as mock
 from osbenchmark.utils import opts
 
 
@@ -141,6 +142,26 @@ def test_make_list_of_close_matches_returns_empty_list_with_no_close_matches(sel
                 [])
         )
 
+    def test_to_dict_str(self):
+        json = '{ "field": 2 }'
+        rsl = opts.to_dict(json)
+        self.assertEqual(rsl, {"field": 2})
+
+        rsl = opts.to_dict('a:1,b:2')
+        self.assertEqual(rsl, {'a': 1, 'b': 2})
+
+    @mock.patch("json.loads")
+    def test_to_dict_file(self, json_loads):
+        mo = mock.mock_open()
+        with mock.patch("builtins.open", mo):
+            opts.to_dict('params.json')
+        mo.assert_called()
+
+        mo = mock.mock_open()
+        with mock.patch("builtins.open", mo):
+            opts.to_dict('index_body:idx.json')
+        mo.assert_not_called()
+
 
 class TestTargetHosts(TestCase):
     def test_empty_arg_parses_as_empty_list(self):