workspace/db.py

import duckdb
import pytest
import os


# 関数の定義部分
def create_table_from_csv(db_path, csv_file_path, table_name):
    conn = duckdb.connect(db_path)
    conn.execute(
        f"""
        CREATE TABLE {table_name} AS
        SELECT * FROM read_csv_auto('{csv_file_path}')
    """
    )
    conn.close()


def fetch_all_from_table(db_path, table_name):
    conn = duckdb.connect(db_path)
    result = conn.execute(f"SELECT * FROM {table_name}").fetchall()
    conn.close()
    return result


# テスト部分
@pytest.fixture
def setup_csv_file(tmpdir):
    # テスト用のCSVファイルを作成
    csv_file_path = tmpdir.join("test_data.csv")
    with open(csv_file_path, "w") as f:
        f.write("id,name\n")
        f.write("1,John Doe\n")
        f.write("2,Jane Smith\n")
    return csv_file_path


@pytest.fixture
def setup_db_file(tmpdir):
    # テスト用のDuckDBファイルを作成
    db_file_path = tmpdir.join("test_db.db")
    return db_file_path


def test_create_table_from_csv(setup_db_file, setup_csv_file):
    db_path = str(setup_db_file)
    csv_file_path = str(setup_csv_file)
    table_name = "test_table"

    # CSVファイルからテーブルを作成
    create_table_from_csv(db_path, csv_file_path, table_name)

    # データベースからテーブルのデータを取得
    result = fetch_all_from_table(db_path, table_name)

    # 期待されるデータ
    expected_result = [(1, "John Doe"), (2, "Jane Smith")]

    # 結果が期待されるデータと一致することを確認
    assert result == expected_result


# テストを実行するためのメイン関数
if __name__ == "__main__":
    pytest.main([__file__])