fn_registry/python/functions/infra/duckdb_upsert_test.py

"""Tests para duckdb_upsert."""

import os
import sys
import tempfile

import duckdb

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from infra.duckdb_upsert import duckdb_upsert  # noqa: E402


def _fresh_db():
    """Crea un .duckdb temporal con tabla t(k PK, a, b) y devuelve su path."""
    fd, path = tempfile.mkstemp(suffix=".duckdb")
    os.close(fd)
    os.remove(path)  # DuckDB crea el archivo limpio.
    con = duckdb.connect(path)
    con.execute("CREATE TABLE t (k INTEGER PRIMARY KEY, a VARCHAR, b VARCHAR)")
    con.close()
    return path


def _select_row(path, k):
    con = duckdb.connect(path, read_only=True)
    try:
        return con.execute("SELECT k, a, b FROM t WHERE k = ?", [k]).fetchone()
    finally:
        con.close()


def test_upsert_fila_nueva_inserta():
    path = _fresh_db()
    try:
        res = duckdb_upsert(
            path, "t", [{"k": 1, "a": "a1", "b": "b1"}], key_cols=["k"]
        )
        assert res == {"status": "ok", "inserted": 1, "updated": 0}
        assert _select_row(path, 1) == (1, "a1", "b1")
    finally:
        os.remove(path)


def test_update_cols_selectivo_no_pisa_columnas_excluidas():
    path = _fresh_db()
    try:
        duckdb_upsert(path, "t", [{"k": 1, "a": "a1", "b": "b1"}], key_cols=["k"])
        # Re-upsert de la misma k cambiando a y b en el dict, pero solo
        # autorizando actualizar 'a'. 'b' debe conservar el valor viejo.
        res = duckdb_upsert(
            path,
            "t",
            [{"k": 1, "a": "a2", "b": "b2"}],
            key_cols=["k"],
            update_cols=["a"],
        )
        assert res == {"status": "ok", "inserted": 0, "updated": 1}
        assert _select_row(path, 1) == (1, "a2", "b1")  # a cambio, b NO
    finally:
        os.remove(path)


def test_update_cols_vacio_do_nothing_no_cambia_existente():
    path = _fresh_db()
    try:
        duckdb_upsert(path, "t", [{"k": 1, "a": "a1", "b": "b1"}], key_cols=["k"])
        res = duckdb_upsert(
            path,
            "t",
            [{"k": 1, "a": "X", "b": "Y"}],
            key_cols=["k"],
            update_cols=[],
        )
        assert res == {"status": "ok", "inserted": 0, "updated": 1}
        assert _select_row(path, 1) == (1, "a1", "b1")  # intacta
    finally:
        os.remove(path)


def test_varias_filas_a_la_vez_mezcla_insert_y_update():
    path = _fresh_db()
    try:
        duckdb_upsert(path, "t", [{"k": 1, "a": "a1", "b": "b1"}], key_cols=["k"])
        res = duckdb_upsert(
            path,
            "t",
            [
                {"k": 1, "a": "a1b", "b": "b1b"},  # update
                {"k": 2, "a": "a2", "b": "b2"},  # insert
                {"k": 3, "a": "a3", "b": "b3"},  # insert
            ],
            key_cols=["k"],
        )
        assert res == {"status": "ok", "inserted": 2, "updated": 1}
        assert _select_row(path, 1) == (1, "a1b", "b1b")
        assert _select_row(path, 2) == (2, "a2", "b2")
        assert _select_row(path, 3) == (3, "a3", "b3")
    finally:
        os.remove(path)


def test_rows_vacio_devuelve_cero():
    path = _fresh_db()
    try:
        res = duckdb_upsert(path, "t", [], key_cols=["k"])
        assert res == {"status": "ok", "inserted": 0, "updated": 0}
    finally:
        os.remove(path)


def test_columnas_inconsistentes_devuelve_error():
    path = _fresh_db()
    try:
        res = duckdb_upsert(
            path,
            "t",
            [{"k": 1, "a": "a1", "b": "b1"}, {"k": 2, "a": "a2"}],
            key_cols=["k"],
        )
        assert res["status"] == "error"
    finally:
        os.remove(path)


def test_identificador_invalido_devuelve_error():
    path = _fresh_db()
    try:
        res = duckdb_upsert(
            path, "t; DROP TABLE t", [{"k": 1, "a": "a1", "b": "b1"}], key_cols=["k"]
        )
        assert res["status"] == "error"
    finally:
        os.remove(path)