dify/api/tests/unit_tests/libs/test_pandas.py

import pandas as pd


def test_pandas_csv(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    data = {"col1": [1, 2.2, -3.3, 4.0, 5], "col2": ["A", "B", "C", "D", "E"]}
    df1 = pd.DataFrame(data)

    # write to csv file
    csv_file_path = tmp_path.joinpath("example.csv")
    df1.to_csv(csv_file_path, index=False)

    # read from csv file
    df2 = pd.read_csv(csv_file_path, on_bad_lines="skip")
    assert df2[df2.columns[0]].to_list() == data["col1"]
    assert df2[df2.columns[1]].to_list() == data["col2"]


def test_pandas_xlsx(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    data = {"col1": [1, 2.2, -3.3, 4.0, 5], "col2": ["A", "B", "C", "D", "E"]}
    df1 = pd.DataFrame(data)

    # write to xlsx file
    xlsx_file_path = tmp_path.joinpath("example.xlsx")
    df1.to_excel(xlsx_file_path, index=False)

    # read from xlsx file
    df2 = pd.read_excel(xlsx_file_path)
    assert df2[df2.columns[0]].to_list() == data["col1"]
    assert df2[df2.columns[1]].to_list() == data["col2"]


def test_pandas_xlsx_with_sheets(tmp_path, monkeypatch):
    monkeypatch.chdir(tmp_path)
    data1 = {"col1": [1, 2, 3, 4, 5], "col2": ["A", "B", "C", "D", "E"]}
    df1 = pd.DataFrame(data1)

    data2 = {"col1": [6, 7, 8, 9, 10], "col2": ["F", "G", "H", "I", "J"]}
    df2 = pd.DataFrame(data2)

    # write to xlsx file with sheets
    xlsx_file_path = tmp_path.joinpath("example_with_sheets.xlsx")
    sheet1 = "Sheet1"
    sheet2 = "Sheet2"
    with pd.ExcelWriter(xlsx_file_path) as excel_writer:
        df1.to_excel(excel_writer, sheet_name=sheet1, index=False)
        df2.to_excel(excel_writer, sheet_name=sheet2, index=False)

    # read from xlsx file with sheets
    with pd.ExcelFile(xlsx_file_path) as excel_file:
        df1 = pd.read_excel(excel_file, sheet_name=sheet1)
        assert df1[df1.columns[0]].to_list() == data1["col1"]
        assert df1[df1.columns[1]].to_list() == data1["col2"]

        df2 = pd.read_excel(excel_file, sheet_name=sheet2)
        assert df2[df2.columns[0]].to_list() == data2["col1"]
        assert df2[df2.columns[1]].to_list() == data2["col2"]