fix test

Author: sfc-gh-yuwang

src/snowflake/snowpark/dataframe_reader.py

@@ -1384,6 +1384,12 @@ def _read_semi_structured_file(self, path: str, format: str) -> DataFrame:
 
         if self._user_schema and format.lower() not in ["json", "xml"]:
             raise ValueError(f"Read {format} does not support user schema")
+        if (
+            self._user_schema
+            and format.lower() == "xml"
+            and XML_ROW_TAG_STRING not in self._cur_options
+        ):
+            raise ValueError("When read XML with user schema, rowtag must be set.")
         path = _validate_stage_path(path)
         self._file_path = path
         self._file_type = format

tests/integ/scala/test_dataframe_reader_suite.py

@@ -1626,10 +1626,6 @@ def test_read_xml_with_no_schema(session, mode, resources_path):
     res = df1.where(sql_expr("xmlget($1, 'num', 0):\"$\"") > 1).collect()
     assert res == [Row("<test>\n  <num>2</num>\n  <str>str2</str>\n</test>")]
 
-    # assert user cannot input a schema to read json
-    with pytest.raises(ValueError):
-        get_reader(session, mode).schema(user_schema).xml(path)
-
     # assert local directory is invalid
     with pytest.raises(
         ValueError, match="DataFrameReader can only read files from stage locations."

tests/integ/test_xml_reader_row_tag.py

@@ -609,3 +609,19 @@ def test_read_xml_with_custom_schema(session):
     ]
     Utils.check_answer(df, expected_result)
     assert df.schema == expected_schema
+
+
+def test_user_schema_without_rowtag(session):
+    user_schema = StructType(
+        [
+            StructField("Author", StringType(), True),
+            StructField("Title", StringType(), True),
+            StructField("genre", StringType(), True),
+            StructField("PRICE", DoubleType(), True),
+            StructField("publish_Date", DateType(), True),
+        ]
+    )
+    with pytest.raises(
+        ValueError, match="When read XML with user schema, rowtag must be set."
+    ):
+        session.read.schema(user_schema).xml(f"@{tmp_stage_name}/{test_file_books_xml}")