from etl_to_crawler.connections import _NulSafePgConnection


class _Cursor:
    def __init__(self):
        self.executed = []
        self.executed_many = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return False

    def execute(self, operation, parameters=None):
        self.executed.append((operation, parameters))

    def executemany(self, operation, parameter_sets):
        self.executed_many.append((operation, parameter_sets))


class _Connection:
    def __init__(self):
        self.raw_cursor = _Cursor()

    def cursor(self):
        return self.raw_cursor


def test_pg_connection_removes_nul_from_execute_parameters():
    raw_connection = _Connection()
    connection = _NulSafePgConnection(raw_connection)

    with connection.cursor() as cur:
        cur.execute('INSERT INTO example VALUES (%s, %s)', ('歌\x00词', {'text': '介\x00绍'}))

    assert raw_connection.raw_cursor.executed == [
        ('INSERT INTO example VALUES (%s, %s)', ('歌词', {'text': '介绍'})),
    ]


def test_pg_connection_removes_nul_from_every_executemany_row():
    raw_connection = _Connection()
    connection = _NulSafePgConnection(raw_connection)

    with connection.cursor() as cur:
        cur.executemany('INSERT INTO example VALUES (%s)', [('A\x00',), ('B\x00',)])

    assert raw_connection.raw_cursor.executed_many == [
        ('INSERT INTO example VALUES (%s)', [('A',), ('B',)]),
    ]