mockgres 0.0.29

An in-memory database that replicates a reasonable subset of Postgres functionality to make unit tests that rely on a database to run.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

#![allow(dead_code)]

use std::net::SocketAddr;
use std::sync::Arc;

use tokio::net::TcpListener;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio_postgres::{Client, NoTls, SimpleQueryMessage};

pub struct TestCtx {
    pub client: Client,
    #[allow(dead_code)]
    conn_str: String,
    pub _bg: JoinHandle<()>,
    pub _server: JoinHandle<()>,
    #[allow(dead_code)]
    pub shutdown: tokio::sync::oneshot::Sender<()>,
    #[allow(dead_code)]
    extra_connections: Arc<Mutex<Vec<JoinHandle<()>>>>,
}

pub async fn start() -> TestCtx {
    let handler = Arc::new(mockgres::Mockgres::default());
    start_with_handler(handler).await
}

pub async fn start_with_handler(handler: Arc<mockgres::Mockgres>) -> TestCtx {
    let (addr, server_task, shutdown) = spawn_server(handler).await;
    let conn_str = format!("host={} port={} user=postgres", addr.ip(), addr.port());
    let (client, connection) = tokio_postgres::connect(&conn_str, NoTls)
        .await
        .expect("connect");
    let conn_task = tokio::spawn(async move {
        if let Err(e) = connection.await {
            eprintln!("connection error: {e}");
        }
    });
    let extras = Arc::new(Mutex::new(Vec::new()));

    TestCtx {
        client,
        conn_str,
        _bg: conn_task,
        _server: server_task,
        shutdown,
        extra_connections: extras,
    }
}

pub async fn spawn_server(
    handler: Arc<mockgres::Mockgres>,
) -> (SocketAddr, JoinHandle<()>, tokio::sync::oneshot::Sender<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
    let addr = listener.local_addr().expect("local addr");
    let (shutdown, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
    let server_task = tokio::spawn(run_server(listener, handler, shutdown_rx));
    (addr, server_task, shutdown)
}

async fn run_server(
    listener: TcpListener,
    handler: Arc<mockgres::Mockgres>,
    mut shutdown: tokio::sync::oneshot::Receiver<()>,
) {
    loop {
        tokio::select! {
            _ = &mut shutdown => break,
            accept_result = listener.accept() => {
                match accept_result {
                    Ok((socket, _peer)) => {
                        let h2 = handler.clone();
                        tokio::spawn(async move {
                            let _ = mockgres::process_socket_with_terminate(socket, None, h2).await;
                        });
                    }
                    Err(e) => {
                        eprintln!("listener accept error: {e}");
                        break;
                    }
                }
            }
        }
    }
}

impl TestCtx {
    #[allow(dead_code)]
    pub async fn new_client(&self) -> Client {
        let (client, connection) = tokio_postgres::connect(&self.conn_str, NoTls)
            .await
            .expect("connect extra client");
        let handle = tokio::spawn(async move {
            if let Err(e) = connection.await {
                eprintln!("connection error: {e}");
            }
        });
        self.extra_connections.lock().await.push(handle);
        client
    }
}

// run a simple query and return the first cell as string
#[allow(dead_code)]
pub async fn simple_first_cell(client: &Client, sql: &str) -> String {
    let msgs = client.simple_query(sql).await.expect("simple query");
    let row = msgs
        .iter()
        .find_map(|m| match m {
            SimpleQueryMessage::Row(r) => Some(r),
            _ => None,
        })
        .expect("expected one row");
    row.get(0).expect("one column").to_string()
}

#[allow(dead_code)]
pub fn assert_db_error_contains(err: &tokio_postgres::Error, needle: &str) {
    let db_err = err.as_db_error().expect("expected db error");
    assert!(
        db_err.message().contains(needle),
        "unexpected error: {:?}",
        db_err.message()
    );
}