Function test_load_generator

src/environmentd/tests/sql.rs:2141–2184

()

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]
fn test_load_generator() {
    let server = test_util::TestHarness::default()
        .unsafe_mode()
        .start_blocking();
    server.enable_feature_flags(&[
        "enable_create_table_from_source",
        "enable_load_generator_counter",
    ]);
    let mut client = server.connect(postgres::NoTls).unwrap();

    client
        .batch_execute("CREATE SOURCE counter FROM LOAD GENERATOR COUNTER (TICK INTERVAL '1ms')")
        .unwrap();

    client
        .batch_execute("CREATE TABLE counter_tbl FROM SOURCE counter")
        .unwrap();

    let row = client
        .query_one("SELECT count(*), mz_now()::text FROM counter_tbl", &[])
        .unwrap();
    let initial_count: i64 = row.get(0);
    let timestamp_millis: String = row.get(1);
    let timestamp_millis: i64 = timestamp_millis.parse().unwrap();
    const WAIT: i64 = 100;
    let next = timestamp_millis + WAIT;
    let expect = initial_count + WAIT;
    Retry::default()
        .retry(|_| {
            let row = client
                .query_one(
                    &format!("SELECT count(*) FROM counter_tbl AS OF AT LEAST {next}"),
                    &[],
                )
                .unwrap();
            let count: i64 = row.get(0);
            if count < expect {
                Err(format!("expected {expect}, got {count}"))
            } else {
                Ok(())
            }
        })
        .unwrap();
}

#[mz_ore::test]
#[allow(clippy::disallowed_methods)]