1- use futures:: { Future , Stream } ;
1+ use futures:: { FutureExt , TryStreamExt } ;
22use native_tls:: { self , Certificate } ;
33use tokio:: net:: TcpStream ;
4- use tokio:: runtime:: current_thread:: Runtime ;
54use tokio_postgres:: tls:: TlsConnect ;
65
76#[ cfg( feature = "runtime" ) ]
87use crate :: MakeTlsConnector ;
98use crate :: TlsConnector ;
109
11- fn smoke_test < T > ( s : & str , tls : T )
10+ async fn smoke_test < T > ( s : & str , tls : T )
1211where
1312 T : TlsConnect < TcpStream > ,
14- T :: Stream : ' static ,
13+ T :: Stream : ' static + Send ,
1514{
16- let mut runtime = Runtime :: new ( ) . unwrap ( ) ;
15+ let stream = TcpStream :: connect ( "127.0.0.1:5433" ) . await . unwrap ( ) ;
1716
1817 let builder = s. parse :: < tokio_postgres:: Config > ( ) . unwrap ( ) ;
18+ let ( mut client, connection) = builder. connect_raw ( stream, tls) . await . unwrap ( ) ;
1919
20- let handshake = TcpStream :: connect ( & "127.0.0.1:5433" . parse ( ) . unwrap ( ) )
21- . map_err ( |e| panic ! ( "{}" , e) )
22- . and_then ( |s| builder. connect_raw ( s, tls) ) ;
23- let ( mut client, connection) = runtime. block_on ( handshake) . unwrap ( ) ;
24- let connection = connection. map_err ( |e| panic ! ( "{}" , e) ) ;
25- runtime. spawn ( connection) ;
20+ let connection = connection. map ( |r| r. unwrap ( ) ) ;
21+ tokio:: spawn ( connection) ;
2622
27- let prepare = client. prepare ( "SELECT 1::INT4" ) ;
28- let statement = runtime. block_on ( prepare) . unwrap ( ) ;
29- let select = client. query ( & statement, & [ ] ) . collect ( ) . map ( |rows| {
30- assert_eq ! ( rows. len( ) , 1 ) ;
31- assert_eq ! ( rows[ 0 ] . get:: <_, i32 >( 0 ) , 1 ) ;
32- } ) ;
33- runtime. block_on ( select) . unwrap ( ) ;
23+ let stmt = client. prepare ( "SELECT $1::INT4" ) . await . unwrap ( ) ;
24+ let rows = client
25+ . query ( & stmt, & [ & 1i32 ] )
26+ . try_collect :: < Vec < _ > > ( )
27+ . await
28+ . unwrap ( ) ;
3429
35- drop ( statement) ;
36- drop ( client) ;
37- runtime. run ( ) . unwrap ( ) ;
30+ assert_eq ! ( rows. len( ) , 1 ) ;
31+ assert_eq ! ( rows[ 0 ] . get:: <_, i32 >( 0 ) , 1 ) ;
3832}
3933
40- #[ test]
41- fn require ( ) {
34+ #[ tokio :: test]
35+ async fn require ( ) {
4236 let connector = native_tls:: TlsConnector :: builder ( )
4337 . add_root_certificate (
4438 Certificate :: from_pem ( include_bytes ! ( "../../test/server.crt" ) ) . unwrap ( ) ,
@@ -48,11 +42,12 @@ fn require() {
4842 smoke_test (
4943 "user=ssl_user dbname=postgres sslmode=require" ,
5044 TlsConnector :: new ( connector, "localhost" ) ,
51- ) ;
45+ )
46+ . await ;
5247}
5348
54- #[ test]
55- fn prefer ( ) {
49+ #[ tokio :: test]
50+ async fn prefer ( ) {
5651 let connector = native_tls:: TlsConnector :: builder ( )
5752 . add_root_certificate (
5853 Certificate :: from_pem ( include_bytes ! ( "../../test/server.crt" ) ) . unwrap ( ) ,
@@ -62,11 +57,12 @@ fn prefer() {
6257 smoke_test (
6358 "user=ssl_user dbname=postgres" ,
6459 TlsConnector :: new ( connector, "localhost" ) ,
65- ) ;
60+ )
61+ . await ;
6662}
6763
68- #[ test]
69- fn scram_user ( ) {
64+ #[ tokio :: test]
65+ async fn scram_user ( ) {
7066 let connector = native_tls:: TlsConnector :: builder ( )
7167 . add_root_certificate (
7268 Certificate :: from_pem ( include_bytes ! ( "../../test/server.crt" ) ) . unwrap ( ) ,
@@ -76,14 +72,13 @@ fn scram_user() {
7672 smoke_test (
7773 "user=scram_user password=password dbname=postgres sslmode=require" ,
7874 TlsConnector :: new ( connector, "localhost" ) ,
79- ) ;
75+ )
76+ . await ;
8077}
8178
82- #[ test]
79+ #[ tokio :: test]
8380#[ cfg( feature = "runtime" ) ]
84- fn runtime ( ) {
85- let mut runtime = Runtime :: new ( ) . unwrap ( ) ;
86-
81+ async fn runtime ( ) {
8782 let connector = native_tls:: TlsConnector :: builder ( )
8883 . add_root_certificate (
8984 Certificate :: from_pem ( include_bytes ! ( "../../test/server.crt" ) ) . unwrap ( ) ,
@@ -92,14 +87,22 @@ fn runtime() {
9287 . unwrap ( ) ;
9388 let connector = MakeTlsConnector :: new ( connector) ;
9489
95- let connect = tokio_postgres:: connect (
90+ let ( mut client , connection ) = tokio_postgres:: connect (
9691 "host=localhost port=5433 user=postgres sslmode=require" ,
9792 connector,
98- ) ;
99- let ( mut client, connection) = runtime. block_on ( connect) . unwrap ( ) ;
100- let connection = connection. map_err ( |e| panic ! ( "{}" , e) ) ;
101- runtime. spawn ( connection) ;
93+ )
94+ . await
95+ . unwrap ( ) ;
96+ let connection = connection. map ( |r| r. unwrap ( ) ) ;
97+ tokio:: spawn ( connection) ;
98+
99+ let stmt = client. prepare ( "SELECT $1::INT4" ) . await . unwrap ( ) ;
100+ let rows = client
101+ . query ( & stmt, & [ & 1i32 ] )
102+ . try_collect :: < Vec < _ > > ( )
103+ . await
104+ . unwrap ( ) ;
102105
103- let execute = client . simple_query ( "SELECT 1" ) . for_each ( |_| Ok ( ( ) ) ) ;
104- runtime . block_on ( execute ) . unwrap ( ) ;
106+ assert_eq ! ( rows . len ( ) , 1 ) ;
107+ assert_eq ! ( rows [ 0 ] . get :: <_ , i32 > ( 0 ) , 1 ) ;
105108}
0 commit comments