rust-postgres/tokio-postgres/tests/test/runtime.rs
Steven Fackler 1f6d9ddc06 Overhaul query cancellation
Multi-host support means we can't simply take the old approach - we need
to know which of the hosts we actually connected to. It's also nice to
move this from the connection to the client since that's what you'd
normally have access to.
2019-01-06 18:03:51 -08:00

101 lines
2.9 KiB
Rust

use futures::Future;
use std::time::{Duration, Instant};
use tokio::runtime::current_thread::Runtime;
use tokio::timer::Delay;
use tokio_postgres::{NoTls, SqlState};
fn smoke_test(s: &str) {
let mut runtime = Runtime::new().unwrap();
let connect = tokio_postgres::connect(s, NoTls);
let (mut client, connection) = runtime.block_on(connect).unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.spawn(connection);
let execute = client.batch_execute("SELECT 1");
runtime.block_on(execute).unwrap();
}
#[test]
#[ignore] // FIXME doesn't work with our docker-based tests :(
fn unix_socket() {
smoke_test("host=/var/run/postgresql port=5433 user=postgres");
}
#[test]
fn tcp() {
smoke_test("host=localhost port=5433 user=postgres")
}
#[test]
fn multiple_hosts_one_port() {
smoke_test("host=foobar.invalid,localhost port=5433 user=postgres");
}
#[test]
fn multiple_hosts_multiple_ports() {
smoke_test("host=foobar.invalid,localhost port=5432,5433 user=postgres");
}
#[test]
fn wrong_port_count() {
let mut runtime = Runtime::new().unwrap();
let f = tokio_postgres::connect("host=localhost port=5433,5433 user=postgres", NoTls);
runtime.block_on(f).err().unwrap();
let f = tokio_postgres::connect(
"host=localhost,localhost,localhost port=5433,5433 user=postgres",
NoTls,
);
runtime.block_on(f).err().unwrap();
}
#[test]
fn target_session_attrs_ok() {
let mut runtime = Runtime::new().unwrap();
let f = tokio_postgres::connect(
"host=localhost port=5433 user=postgres target_session_attrs=read-write",
NoTls,
);
runtime.block_on(f).unwrap();
}
#[test]
fn target_session_attrs_err() {
let mut runtime = Runtime::new().unwrap();
let f = tokio_postgres::connect(
"host=localhost port=5433 user=postgres target_session_attrs=read-write
options='-c default_transaction_read_only=on'",
NoTls,
);
runtime.block_on(f).err().unwrap();
}
#[test]
fn cancel_query() {
let mut runtime = Runtime::new().unwrap();
let connect = tokio_postgres::connect("host=localhost port=5433 user=postgres", NoTls);
let (mut client, connection) = runtime.block_on(connect).unwrap();
let connection = connection.map_err(|e| panic!("{}", e));
runtime.spawn(connection);
let sleep = client
.batch_execute("SELECT pg_sleep(100)")
.then(|r| match r {
Ok(_) => panic!("unexpected success"),
Err(ref e) if e.code() == Some(&SqlState::QUERY_CANCELED) => Ok::<(), ()>(()),
Err(e) => panic!("unexpected error {}", e),
});
let cancel = Delay::new(Instant::now() + Duration::from_millis(100))
.then(|r| {
r.unwrap();
client.cancel_query(NoTls)
})
.then(|r| {
r.unwrap();
Ok::<(), ()>(())
});
let ((), ()) = runtime.block_on(sleep.join(cancel)).unwrap();
}