sqlite-utils/tests/test_export.py

import csv
import io
import sqlite3

import pytest

from sqlite_utils.export import rows_to_csv_file


def _cursor():
    db = sqlite3.connect(":memory:")
    db.execute("create table creatures (id integer, name text)")
    db.executemany(
        "insert into creatures (id, name) values (?, ?)",
        ((1, "Cleo"), (2, "Cardi, Jr.")),
    )
    return db, db.execute("select id, name from creatures order by id")


def test_rows_to_csv_file():
    db, cursor = _cursor()
    output = io.StringIO(newline="")
    rows_to_csv_file(cursor, output, lineterminator="\n")
    assert output.getvalue() == 'id,name\n1,Cleo\n2,"Cardi, Jr."\n'
    db.close()


def test_rows_to_csv_file_without_header():
    db, cursor = _cursor()
    output = io.StringIO(newline="")
    rows_to_csv_file(cursor, output, header=False, lineterminator="\n")
    assert output.getvalue() == '1,Cleo\n2,"Cardi, Jr."\n'
    db.close()


def test_rows_to_csv_file_tsv():
    db, cursor = _cursor()
    output = io.StringIO(newline="")
    rows_to_csv_file(cursor, output, dialect="excel-tab", lineterminator="\n")
    assert output.getvalue() == "id\tname\n1\tCleo\n2\tCardi, Jr.\n"
    db.close()


def test_rows_to_csv_file_accepts_dialect_class():
    class PipeDialect(csv.excel):
        delimiter = "|"

    db, cursor = _cursor()
    output = io.StringIO(newline="")
    rows_to_csv_file(cursor, output, dialect=PipeDialect, lineterminator="\n")
    assert output.getvalue() == "id|name\n1|Cleo\n2|Cardi, Jr.\n"
    db.close()


def test_rows_to_csv_file_with_mapping_row_factory():
    db = sqlite3.connect(":memory:")
    db.execute("create table creatures (id integer, name text)")
    db.executemany(
        "insert into creatures (id, name) values (?, ?)",
        ((1, "Cleo"), (2, "Cardi, Jr.")),
    )

    def dict_factory(cursor, row):
        return {
            column[0]: value
            for column, value in zip(cursor.description, row)
        }

    db.row_factory = dict_factory
    cursor = db.execute("select id, name from creatures order by id")
    output = io.StringIO(newline="")
    rows_to_csv_file(cursor, output, lineterminator="\n")
    assert output.getvalue() == 'id,name\n1,Cleo\n2,"Cardi, Jr."\n'
    db.close()


@pytest.mark.parametrize("header", (True, False))
def test_rows_to_csv_file_requires_result_columns(header):
    db = sqlite3.connect(":memory:")
    cursor = db.execute("create table creatures (id integer)")
    output = io.StringIO(newline="")
    with pytest.raises(ValueError, match="Cursor does not have result columns"):
        rows_to_csv_file(cursor, output, header=header)
    assert output.getvalue() == ""
    db.close()