From 6d08aaa6bacbe64ce8a359fd330d1d2034114157 Mon Sep 17 00:00:00 2001
From: Grieve <grieve54706@gmail.com>
Date: Tue, 3 Dec 2024 11:53:46 +0800
Subject: [PATCH] fix: wren core raise CoreError instead of
 pyo3_runtime.PanicException

---
 wren-core-py/src/extractor.rs            |  6 +++---
 wren-core-py/tests/test_modeling_core.py | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/wren-core-py/src/extractor.rs b/wren-core-py/src/extractor.rs
index 3636ba5ee..26f8c4baa 100644
--- a/wren-core-py/src/extractor.rs
+++ b/wren-core-py/src/extractor.rs
@@ -16,10 +16,10 @@ pub struct PyExtractor {
 #[pymethods]
 impl PyExtractor {
     #[new]
-    pub fn new(mdl_base64: &str) -> Self {
-        let manifest = to_manifest(mdl_base64).unwrap();
+    pub fn new(mdl_base64: &str) -> Result<Self, CoreError> {
+        let manifest = to_manifest(mdl_base64)?;
         let mdl = WrenMDL::new_ref(manifest);
-        Self { mdl }
+        Ok(Self { mdl })
     }
 
     /// parse the given SQL and return the list of used table name.
diff --git a/wren-core-py/tests/test_modeling_core.py b/wren-core-py/tests/test_modeling_core.py
index 5b162e6b5..814a679a8 100644
--- a/wren-core-py/tests/test_modeling_core.py
+++ b/wren-core-py/tests/test_modeling_core.py
@@ -142,6 +142,29 @@ def test_get_available_functions():
     assert max_if["param_types"] is None
 
 
+@pytest.mark.parametrize(
+    ("value", "expected_error", "error_message"),
+    [
+        (
+            None,
+            TypeError,
+            "argument 'mdl_base64': 'NoneType' object cannot be converted to 'PyString'",
+        ),
+        ("xxx", Exception, "Base64 decode error: Invalid padding"),
+        ("{}", Exception, "Base64 decode error: Invalid symbol 123, offset 0."),
+        (
+            "",
+            Exception,
+            "Serde JSON error: EOF while parsing a value at line 1 column 0",
+        ),
+    ],
+)
+def test_extractor_with_invalid_manifest(value, expected_error, error_message):
+    with pytest.raises(expected_error) as e:
+        Extractor(value)
+    assert str(e.value) == error_message
+
+
 @pytest.mark.parametrize(
     ("sql", "expected"),
     [