import argparse

import psycopg2
from psycopg2.extras import wait_select


parser = argparse.ArgumentParser()
parser.add_argument('connstr', help='libpq connection string')
parser.add_argument('--num', type=int, default=1000, help='number of notifies')

args = parser.parse_args()

conn = psycopg2.connect(args.connstr, async=True)
wait_select(conn)
cur = conn.cursor()

cur.execute('''
CREATE OR REPLACE FUNCTION notify_test(count INT)
RETURNS VOID AS
$$
BEGIN
    FOR i IN 1..count LOOP
        PERFORM pg_notify('notify_test', i::TEXT);
    END LOOP;
END;
$$
LANGUAGE PLPGSQL;
''')
wait_select(conn)

cur.execute('LISTEN notify_test')
wait_select(conn)

cur.execute('SELECT notify_test(%s)', (args.num, ))
wait_select(conn)

print 'got', len(conn.notifies), 'notifies, expected', args.num