Skip to content

Commit 8acd292

Browse files
culpgrantsfc-gh-mmishchenko
authored andcommitted
move to a pure unit test with full mocking
1 parent d0ccb85 commit 8acd292

File tree

2 files changed

+45
-41
lines changed

2 files changed

+45
-41
lines changed

test/integ/pandas/test_pandas_tools.py

-41
Original file line numberDiff line numberDiff line change
@@ -483,47 +483,6 @@ def test_write_pandas_use_logical_type(
483483
cnx.execute_string(drop_sql)
484484

485485

486-
@pytest.mark.parametrize(
487-
("use_vectorized_scanner", "expected_file_format"),
488-
[
489-
(None, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto)"),
490-
(True, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto USE_VECTORIZED_SCANNER=TRUE)"),
491-
(False, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto USE_VECTORIZED_SCANNER=FALSE)"),
492-
],
493-
)
494-
def test_write_pandas_use_vectorized_scanner(
495-
conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
496-
use_vectorized_scanner: bool | None,
497-
expected_file_format: str,
498-
):
499-
"""Test that use_vectorized_scanner is making correct arguments to the COPY INTO command in SQL."""
500-
from snowflake.connector.cursor import SnowflakeCursor
501-
502-
table_name = random_string(5, "use_vectorized_scanner")
503-
504-
with conn_cnx() as cnx:
505-
def mocked_execute(*args, **kwargs):
506-
if len(args) >= 1 and args[0].startswith("COPY INTO"):
507-
assert expected_file_format in args[0]
508-
cur = SnowflakeCursor(cnx)
509-
cur._result = iter([])
510-
return cur
511-
512-
with mock.patch(
513-
"snowflake.connector.cursor.SnowflakeCursor.execute",
514-
side_effect=mocked_execute,
515-
) as m_execute:
516-
success, nchunks, nrows, _ = write_pandas(
517-
cnx,
518-
sf_connector_version_df.get(),
519-
table_name=table_name,
520-
use_vectorized_scanner=use_vectorized_scanner,
521-
)
522-
assert m_execute.called and any(
523-
map(lambda e: "COPY INTO" in str(e[0]), m_execute.call_args_list)
524-
)
525-
526-
527486
def test_invalid_table_type_write_pandas(
528487
conn_cnx: Callable[..., Generator[SnowflakeConnection]],
529488
):

test/unit/test_pandas_tools.py

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
from unittest.mock import MagicMock
2+
from typing import Union
3+
4+
import pandas as pd
5+
import pytest
6+
7+
from .mock_utils import mock_connection
8+
from snowflake.connector import pandas_tools
9+
10+
11+
@pytest.mark.parametrize(
12+
("use_vectorized_scanner", "expected_file_format"),
13+
[
14+
(None, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto)"),
15+
(True, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto USE_VECTORIZED_SCANNER=TRUE)"),
16+
(False, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto USE_VECTORIZED_SCANNER=FALSE)"),
17+
],
18+
)
19+
def test_write_pandas_use_vectorized_scanner(use_vectorized_scanner: Union[bool, None], expected_file_format: str):
20+
# Setup Mocks
21+
df = pd.DataFrame({"col1": [1, 2, 3]})
22+
23+
mock_conn = mock_connection()
24+
mock_cursor = MagicMock()
25+
mock_conn.cursor.return_value = mock_cursor
26+
27+
# Execute Function
28+
pandas_tools.write_pandas(
29+
conn=mock_conn,
30+
df=df,
31+
table_name="test_table",
32+
schema="test_schema",
33+
database="test_database",
34+
use_vectorized_scanner=use_vectorized_scanner,
35+
)
36+
37+
executed_sql_statements = [call[0][0] for call in mock_cursor.execute.call_args_list]
38+
39+
assert any(
40+
'COPY INTO' in sql and expected_file_format in sql
41+
for sql in executed_sql_statements
42+
)
43+
44+
45+

0 commit comments

Comments
 (0)