Merge pull request #74 from yassun7010/add_testcase

test: add test case.

Author: ya7010

turu-snowflake/tests/turu/test_snowflake.py

@@ -244,22 +244,47 @@ def test_record_pandas_dataframe(self, connection: turu.snowflake.Connection):
         import pandas as pd  # type: ignore[import]
         from pandas.testing import assert_frame_equal  # type: ignore[import]
 
+        with tempfile.NamedTemporaryFile() as file:
+            with record_to_csv(
+                file.name,
+                connection.execute_map(
+                    pd.DataFrame, "select 1 as ID union all select 2 AS ID"
+                ),
+            ) as cursor:
+                expected = pd.DataFrame({"ID": [1, 2]}, dtype="int8")
+
+                assert_frame_equal(cursor.fetch_pandas_all(), expected)
+
+            assert (
+                Path(file.name).read_text()
+                == dedent(
+                    """
+                    ID
+                    1
+                    2
+                    """
+                ).lstrip()
+            )
+
+    @pytest.mark.skipif(not USE_PANDAS, reason="pandas is not installed")
+    def test_record_pandas_dataframe_with_limit_option(
+        self, connection: turu.snowflake.Connection
+    ):
+        import pandas as pd  # type: ignore[import]
+        from pandas.testing import assert_frame_equal  # type: ignore[import]
+
         with tempfile.NamedTemporaryFile() as file:
             with record_to_csv(
                 file.name,
                 connection.execute_map(
                     pd.DataFrame,
-                    "select 1 as ID union all select 2 AS ID",
+                    "select value::integer as ID from table(flatten(ARRAY_GENERATE_RANGE(1, 10)))",
                 ),
+                limit=2,
             ) as cursor:
-                expected = pd.DataFrame(
-                    {"ID": [1, 2]},
-                    dtype="int8",
-                )
+                expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object")
 
                 assert_frame_equal(cursor.fetch_pandas_all(), expected)
-                for row in expected.values:
-                    print(row)
 
             assert (
                 Path(file.name).read_text()