from contextlib import closing

from fixtures.zenith_fixtures import ZenithEnvBuilder
from fixtures.benchmark_fixture import ZenithBenchmarker


def test_startup(zenith_env_builder: ZenithEnvBuilder, zenbenchmark: ZenithBenchmarker):
    zenith_env_builder.num_safekeepers = 3
    env = zenith_env_builder.init_start()

    # Start
    env.zenith_cli.create_branch('test_startup')
    with zenbenchmark.record_duration("startup_time"):
        pg = env.postgres.create_start('test_startup')
        pg.safe_psql("select 1;")

    # Restart
    pg.stop_and_destroy()
    with zenbenchmark.record_duration("restart_time"):
        pg.create_start('test_startup')
        pg.safe_psql("select 1;")

    # Fill up
    num_rows = 1000000  # 30 MB
    num_tables = 100
    with closing(pg.connect()) as conn:
        with conn.cursor() as cur:
            for i in range(num_tables):
                cur.execute(f'create table t_{i} (i integer);')
                cur.execute(f'insert into t_{i} values (generate_series(1,{num_rows}));')

    # Read
    with zenbenchmark.record_duration("read_time"):
        pg.safe_psql("select * from t_0;")

    # Read again
    with zenbenchmark.record_duration("second_read_time"):
        pg.safe_psql("select * from t_0;")

    # Restart
    pg.stop_and_destroy()
    with zenbenchmark.record_duration("restart_with_data"):
        pg.create_start('test_startup')
        pg.safe_psql("select 1;")

    # Read
    with zenbenchmark.record_duration("read_after_restart"):
        pg.safe_psql("select * from t_0;")