update exceptiuon

kgyrtkirk · May 7, 2024 · 2051dbc · 2051dbc
1 parent c46a65c
commit 2051dbc
Showing 1 changed file with 11 additions and 6 deletions.
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/SqlTestFrameworkConfig.java b/sql/src/test/java/org/apache/druid/sql/calcite/SqlTestFrameworkConfig.java
@@ -145,7 +145,8 @@ public SqlTestFrameworkConfig(Map<String, String> queryParams)
   @Nonnull
   private <T> T getValueFromMap(Map<String, String> map, Class<? extends Annotation> annotationClass) throws Exception
   {
-    String value = map.get(annotationClass.getSimpleName());
+    String key = annotationClass.getSimpleName();
+    String value = map.get(key);
     if (value == null) {
       return defaultValue(annotationClass);
     }
@@ -157,7 +158,7 @@ private <T> T getValueFromMap(Map<String, String> map, Class<? extends Annotatio
     if (type == Class.class) {
       return (T) getQueryComponentSupplierForName(value);
     }
-    throw new RuntimeException("don't know how to handle conversion to " + type);
+    throw new IAE("Cannot handle conversion of key [%s] with value [%s] to type [%s].", key, value, type);
   }
 
   static LoadingCache<String, Set<Class<? extends QueryComponentSupplier>>> componentSupplierClassCache = CacheBuilder
@@ -309,7 +310,7 @@ public SqlTestFrameworkConfig getConfig()
 
     public SqlTestFramework get() throws Exception
     {
-      return configStore.getConfigurationInstance(config, x -> x).framework;
+      return configStore.getConfigurationInstance(config, null).framework;
     }
 
     public <T extends Annotation> T getAnnotation(Class<T> annotationType)
@@ -342,14 +343,18 @@ public ConfigurationInstance(
         SqlTestFrameworkConfig config,
         Function<QueryComponentSupplier, QueryComponentSupplier> queryComponentSupplierWrapper) throws Exception
     {
-      this(config, queryComponentSupplierWrapper.apply(makeQueryComponentSupplier(config.componentSupplier)));
+      this(config, makeQueryComponentSupplier(config.componentSupplier, queryComponentSupplierWrapper));
     }
 
     private static QueryComponentSupplier makeQueryComponentSupplier(
-        Class<? extends QueryComponentSupplier> supplierClazz) throws Exception
+        Class<? extends QueryComponentSupplier> supplierClazz, Function<QueryComponentSupplier, QueryComponentSupplier> wrapper) throws Exception
     {
       Constructor<? extends QueryComponentSupplier> constructor = supplierClazz.getConstructor(TempDirProducer.class);
-      return constructor.newInstance(new TempDirProducer("druid-test"));
+      QueryComponentSupplier componentSupplier = constructor.newInstance(new TempDirProducer("druid-test"));
+      if (wrapper != null) {
+        return wrapper.apply(componentSupplier);
+      }
+      return componentSupplier;
     }
 
     public void close()