<reponame>Zeto-Ltd/replisome
import datetime

from replisome.receivers import Wal2JsonReceiver


def create_alarm(cursor, severity=50, cause="MAX", time_created=None):
    time_created = time_created or datetime.datetime(
        year=2021, month=6, day=17, hour=11
    )
    insert_args = {
        "severity": severity,
        "cause": cause,
        "time_created": time_created,
    }
    cursor.execute(
        """
INSERT INTO alarm (severity, cause, time_created)
VALUES (%(severity)s, %(cause)s, %(time_created)s)
RETURNING alarm_id;
    """,
        insert_args,
    )
    insert_args["alarm_id"] = cursor.fetchone()[0]
    return insert_args


def update_alarm(
    cursor, alarm_id, severity=50, cause="MAX", time_created=None
):
    time_created = time_created or datetime.datetime(
        year=2021, month=6, day=17, hour=11
    )
    update_args = {
        "alarm_id": alarm_id,
        "severity": severity,
        "cause": cause,
        "time_created": time_created,
    }
    cursor.execute(
        """
UPDATE alarm
SET severity = %(severity)s, cause = %(cause)s, time_created = %(time_created)s
WHERE alarm_id = %(alarm_id)s
    """,
        update_args,
    )
    return update_args


def delete_alarm(cursor, alarm_id):
    delete_args = {
        "alarm_id": alarm_id,
    }
    cursor.execute(
        """
DELETE FROM alarm WHERE alarm_id = %(alarm_id)s
    """,
        delete_args,
    )
    return delete_args


def test_can_capture_inserts(test_db_connection):
    changes = []

    def on_callback(change):
        changes.append(change)

    receiver = Wal2JsonReceiver(
        dsn=test_db_connection.dsn,
        slot="wal2json_receiver_insert",
        block=False,
        message_cb=on_callback,
        plugin="wal2json",
    )
    try:
        receiver.start()
        receiver.on_loop(wait_time=0.2)
        with test_db_connection.cursor() as cursor:
            created_alarms = [
                create_alarm(cursor),
                create_alarm(
                    cursor,
                    time_created=datetime.datetime(
                        year=2021, month=6, day=17, hour=12
                    ),
                ),
            ]
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        assert len(changes) == len(created_alarms)
        assert changes == [
            {
                "kind": "insert",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[0]["alarm_id"],
                    50,
                    "MAX",
                    "2021-06-17 11:00:00+00",
                    None,
                    None,
                ],
            },
            {
                "kind": "insert",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[1]["alarm_id"],
                    50,
                    "MAX",
                    "2021-06-17 12:00:00+00",
                    None,
                    None,
                ],
            },
        ]
    finally:
        receiver.close()


def test_can_capture_updates(test_db_connection):
    changes = []

    def on_callback(change):
        changes.append(change)

    receiver = Wal2JsonReceiver(
        dsn=test_db_connection.dsn,
        slot="wal2json_receiver_update",
        block=False,
        message_cb=on_callback,
        plugin="wal2json",
    )
    try:
        with test_db_connection.cursor() as cursor:
            created_alarms = [
                create_alarm(cursor),
                create_alarm(
                    cursor,
                    time_created=datetime.datetime(
                        year=2021, month=6, day=17, hour=12
                    ),
                ),
            ]
        receiver.start()
        receiver.on_loop(wait_time=0.2)
        with test_db_connection.cursor() as cursor:
            updated_alarms = [
                update_alarm(
                    cursor, created_alarms[0]["alarm_id"], severity=30
                ),
                update_alarm(
                    cursor, created_alarms[1]["alarm_id"], severity=70
                ),
            ]
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        assert len(changes) == len(updated_alarms)
        assert changes == [
            {
                "kind": "update",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[0]["alarm_id"],
                    30,
                    "MAX",
                    "2021-06-17 11:00:00+00",
                    None,
                    None,
                ],
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[0]["alarm_id"]],
                },
            },
            {
                "kind": "update",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[1]["alarm_id"],
                    70,
                    "MAX",
                    "2021-06-17 11:00:00+00",
                    None,
                    None,
                ],
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[1]["alarm_id"]],
                },
            },
        ]
    finally:
        receiver.close()


def test_can_capture_deletes(test_db_connection):
    changes = []

    def on_callback(change):
        changes.append(change)

    receiver = Wal2JsonReceiver(
        dsn=test_db_connection.dsn,
        slot="wal2json_receiver_delete",
        block=False,
        message_cb=on_callback,
        plugin="wal2json",
    )
    try:
        with test_db_connection.cursor() as cursor:
            created_alarms = [
                create_alarm(cursor),
                create_alarm(
                    cursor,
                    time_created=datetime.datetime(
                        year=2021, month=6, day=17, hour=12
                    ),
                ),
            ]
        receiver.start()
        receiver.on_loop(wait_time=0.2)
        with test_db_connection.cursor() as cursor:
            deleted_alarms = [
                delete_alarm(cursor, created_alarms[0]["alarm_id"]),
                delete_alarm(cursor, created_alarms[1]["alarm_id"]),
            ]
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        receiver.on_loop(wait_time=0.2)
        assert len(changes) == len(deleted_alarms)
        assert changes == [
            {
                "kind": "delete",
                "schema": "public",
                "table": "alarm",
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[0]["alarm_id"]],
                },
            },
            {
                "kind": "delete",
                "schema": "public",
                "table": "alarm",
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[1]["alarm_id"]],
                },
            },
        ]
    finally:
        receiver.close()


def test_can_capture_chunked_changes(test_db_connection):
    changes = []

    def on_callback(change):
        changes.append(change)

    receiver = Wal2JsonReceiver(
        dsn=test_db_connection.dsn,
        slot="wal2json_receiver_chunked_changes",
        block=False,
        message_cb=on_callback,
        plugin="wal2json",
        options=[["write-in-chunks", "t"]],
    )
    try:
        receiver.start()
        receiver.on_loop(wait_time=0.2)
        with test_db_connection.cursor() as cursor:
            created_alarms = [
                create_alarm(cursor),
                create_alarm(
                    cursor,
                    time_created=datetime.datetime(
                        year=2021, month=6, day=17, hour=12
                    ),
                ),
            ]
            update_alarm(cursor, created_alarms[0]["alarm_id"], severity=30)
            update_alarm(cursor, created_alarms[1]["alarm_id"], severity=70)
            delete_alarm(cursor, created_alarms[0]["alarm_id"])
            delete_alarm(cursor, created_alarms[1]["alarm_id"])
        for index in range(30):
            receiver.on_loop(wait_time=0.2)
        assert len(changes) == 6
        assert changes == [
            {
                "kind": "insert",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[0]["alarm_id"],
                    50,
                    "MAX",
                    "2021-06-17 11:00:00+00",
                    None,
                    None,
                ],
            },
            {
                "kind": "insert",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[1]["alarm_id"],
                    50,
                    "MAX",
                    "2021-06-17 12:00:00+00",
                    None,
                    None,
                ],
            },
            {
                "kind": "update",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[0]["alarm_id"],
                    30,
                    "MAX",
                    "2021-06-17 11:00:00+00",
                    None,
                    None,
                ],
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[0]["alarm_id"]],
                },
            },
            {
                "kind": "update",
                "schema": "public",
                "table": "alarm",
                "columnnames": [
                    "alarm_id",
                    "severity",
                    "cause",
                    "time_created",
                    "time_cleared",
                    "time_deleted",
                ],
                "columntypes": [
                    "bigint",
                    "integer",
                    "character varying(20)",
                    "timestamp with time zone",
                    "timestamp with time zone",
                    "timestamp with time zone",
                ],
                "columnvalues": [
                    created_alarms[1]["alarm_id"],
                    70,
                    "MAX",
                    "2021-06-17 11:00:00+00",
                    None,
                    None,
                ],
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[1]["alarm_id"]],
                },
            },
            {
                "kind": "delete",
                "schema": "public",
                "table": "alarm",
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[0]["alarm_id"]],
                },
            },
            {
                "kind": "delete",
                "schema": "public",
                "table": "alarm",
                "oldkeys": {
                    "keynames": ["alarm_id"],
                    "keytypes": ["bigint"],
                    "keyvalues": [created_alarms[1]["alarm_id"]],
                },
            },
        ]
    finally:
        receiver.close()
