watermelon_testkit/
server.rs

1use std::{
2    env, fmt,
3    net::SocketAddr,
4    path::{Path, PathBuf},
5    process::Stdio,
6    thread,
7    time::Duration,
8};
9
10use hyper::StatusCode;
11use rustix::process::{Pid, Signal, kill_process};
12use serde_json::Value;
13use tempfile::TempDir;
14use tokio::{
15    fs,
16    process::{Child, Command},
17    time::{Instant, sleep},
18};
19
20use crate::http::HttpClient;
21
22const STARTUP_TIMEOUT: Duration = Duration::from_secs(15);
23const LOG_TAIL_LINES: usize = 50;
24const NOT_FOUND: &str =
25    "nats-server binary not found; set NATS_SERVER_BIN or add nats-server to PATH";
26
27/// Configuration for a [`NatsServer`]
28///
29/// Obtained from [`NatsServer::builder`].
30#[derive(Debug, Default)]
31#[must_use]
32pub struct NatsServerBuilder {
33    jetstream: bool,
34    jetstream_domain: Option<String>,
35    websocket: bool,
36    users: Vec<(String, String)>,
37    max_payload: Option<u32>,
38    extra_config: Vec<String>,
39}
40
41/// A `nats-server` process owned by a single test
42///
43/// The process is killed when this is dropped. If the test is panicking,
44/// the tail of the server log is printed and the temporary directory
45/// holding the config, the full log and the `JetStream` store is kept
46/// around for inspection.
47#[derive(Debug)]
48pub struct NatsServer {
49    child: Child,
50    pid: Pid,
51    dir: TempDir,
52    client_url: String,
53    websocket_url: Option<String>,
54    monitoring_addr: SocketAddr,
55    http: HttpClient,
56}
57
58impl NatsServerBuilder {
59    /// Enable `JetStream`, storing data in the server's temporary directory
60    pub fn jetstream(mut self) -> Self {
61        self.jetstream = true;
62        self
63    }
64
65    /// Enable `JetStream` under the given `domain`
66    pub fn jetstream_domain(mut self, domain: impl Into<String>) -> Self {
67        self.jetstream = true;
68        self.jetstream_domain = Some(domain.into());
69        self
70    }
71
72    /// Enable a plaintext WebSocket listener
73    pub fn websocket(mut self) -> Self {
74        self.websocket = true;
75        self
76    }
77
78    /// Require clients to authenticate, allowing `username` with `password`
79    ///
80    /// Can be called multiple times to allow multiple users.
81    pub fn user(mut self, username: impl Into<String>, password: impl Into<String>) -> Self {
82        self.users.push((username.into(), password.into()));
83        self
84    }
85
86    /// Set the maximum payload size accepted by the server
87    pub fn max_payload(mut self, max_payload: u32) -> Self {
88        self.max_payload = Some(max_payload);
89        self
90    }
91
92    /// Append raw lines to the generated server configuration file
93    pub fn config(mut self, config: impl Into<String>) -> Self {
94        self.extra_config.push(config.into());
95        self
96    }
97
98    /// Start the server, waiting until it's ready to accept clients
99    ///
100    /// Returns `None` if no `nats-server` binary could be found
101    /// and the `CI` environment variable isn't set.
102    ///
103    /// # Panics
104    ///
105    /// Panics if the binary is missing while running in CI,
106    /// or if the server fails to start.
107    #[must_use]
108    pub async fn start(self) -> Option<NatsServer> {
109        let Some(bin) = find_nats_server() else {
110            assert!(
111                env::var_os("CI").is_none_or(|ci| ci.is_empty()),
112                "{NOT_FOUND}"
113            );
114            eprintln!("SKIPPED: {NOT_FOUND}");
115            return None;
116        };
117
118        let dir = tempfile::Builder::new()
119            .prefix("watermelon-testkit-")
120            .tempdir()
121            .expect("create server temporary directory");
122        let config_path = dir.path().join("nats-server.conf");
123        fs::write(&config_path, self.render(dir.path()))
124            .await
125            .expect("write server config");
126
127        let child = Command::new(&bin)
128            .arg("--config")
129            .arg(&config_path)
130            .stdin(Stdio::null())
131            .stdout(Stdio::null())
132            .stderr(Stdio::null())
133            .kill_on_drop(true)
134            .spawn()
135            .unwrap_or_else(|err| panic!("spawn {}: {err}", bin.display()));
136        let pid = child
137            .id()
138            .and_then(|id| Pid::from_raw(id.try_into().ok()?))
139            .expect("nats-server pid");
140
141        let mut server = NatsServer {
142            child,
143            pid,
144            dir,
145            client_url: String::new(),
146            websocket_url: None,
147            monitoring_addr: SocketAddr::from(([0, 0, 0, 0], 0)),
148            http: HttpClient::new(),
149        };
150        server.wait_ready(self.jetstream).await;
151        Some(server)
152    }
153
154    fn render(&self, dir: &Path) -> String {
155        let mut config = String::new();
156        self.write_config(&mut config, dir)
157            .expect("formatting into a String never fails");
158        config
159    }
160
161    fn write_config(&self, w: &mut impl fmt::Write, dir: &Path) -> fmt::Result {
162        writeln!(w, r#"listen: "127.0.0.1:-1""#)?;
163        writeln!(w, r#"http: "127.0.0.1:-1""#)?;
164        writeln!(w, "ports_file_dir: {}", quote(dir))?;
165        writeln!(w, "log_file: {}", quote(dir.join("nats-server.log")))?;
166        writeln!(w, "logtime: true")?;
167        writeln!(w, "debug: true")?;
168        writeln!(w, "trace: true")?;
169        writeln!(w, "max_traced_msg_len: 256")?;
170
171        if self.jetstream {
172            write!(
173                w,
174                "jetstream {{ store_dir: {}",
175                quote(dir.join("jetstream"))
176            )?;
177            if let Some(domain) = &self.jetstream_domain {
178                write!(w, ", domain: {}", quote(domain))?;
179            }
180            writeln!(w, " }}")?;
181        }
182        if self.websocket {
183            writeln!(w, r#"websocket {{ listen: "127.0.0.1:-1", no_tls: true }}"#)?;
184        }
185        if !self.users.is_empty() {
186            write!(w, "authorization {{ users: [")?;
187            for (username, password) in &self.users {
188                write!(
189                    w,
190                    " {{ user: {}, password: {} }}",
191                    quote(username),
192                    quote(password)
193                )?;
194            }
195            writeln!(w, " ] }}")?;
196        }
197        if let Some(max_payload) = self.max_payload {
198            writeln!(w, "max_payload: {max_payload}")?;
199        }
200        for extra in &self.extra_config {
201            writeln!(w, "{extra}")?;
202        }
203
204        Ok(())
205    }
206}
207
208impl NatsServer {
209    /// Configure a new server
210    pub fn builder() -> NatsServerBuilder {
211        NatsServerBuilder::default()
212    }
213
214    /// The `nats://` URL clients should connect to
215    #[must_use]
216    pub fn client_url(&self) -> &str {
217        &self.client_url
218    }
219
220    /// The `ws://` URL clients should connect to
221    ///
222    /// Returns `None` unless [`NatsServerBuilder::websocket`] was called.
223    #[must_use]
224    pub fn websocket_url(&self) -> Option<&str> {
225        self.websocket_url.as_deref()
226    }
227
228    /// Query a monitoring endpoint, like `/connz?subs=1` or `/varz`
229    ///
230    /// # Panics
231    ///
232    /// Panics if the request fails or the response isn't successful JSON.
233    #[must_use]
234    pub async fn monitor(&self, path_and_query: &str) -> Value {
235        let (status, body) = self
236            .http
237            .get(self.monitoring_addr, path_and_query)
238            .await
239            .unwrap_or_else(|err| panic!("GET {path_and_query}: {err}"));
240        assert_eq!(
241            status,
242            StatusCode::OK,
243            "GET {path_and_query}: {}",
244            String::from_utf8_lossy(&body)
245        );
246        serde_json::from_slice(&body)
247            .unwrap_or_else(|err| panic!("GET {path_and_query}: invalid JSON: {err}"))
248    }
249
250    /// The connections currently open with the server, including their subscriptions
251    ///
252    /// Shorthand for the `connections` array of `/connz?subs=1`.
253    #[must_use]
254    pub async fn connections(&self) -> Vec<Value> {
255        match self
256            .monitor("/connz?subs=1")
257            .await
258            .get_mut("connections")
259            .map(Value::take)
260        {
261            Some(Value::Array(connections)) => connections,
262            _ => Vec::new(),
263        }
264    }
265
266    /// Put the server into Lame Duck Mode
267    pub fn lame_duck(&self) {
268        self.signal(Signal::USR2);
269    }
270
271    /// Freeze the server process, leaving its sockets open
272    pub fn pause(&self) {
273        self.signal(Signal::STOP);
274    }
275
276    /// Resume a server frozen by [`NatsServer::pause`]
277    pub fn resume(&self) {
278        self.signal(Signal::CONT);
279    }
280
281    /// Gracefully shut the server down and wait for it to exit
282    ///
283    /// # Panics
284    ///
285    /// Panics if waiting on the process fails.
286    pub async fn shutdown(&mut self) {
287        self.signal(Signal::TERM);
288        self.child
289            .wait()
290            .await
291            .expect("wait for nats-server to exit");
292    }
293
294    /// The full server log
295    pub async fn log(&self) -> String {
296        fs::read_to_string(self.log_path())
297            .await
298            .unwrap_or_default()
299    }
300
301    fn log_path(&self) -> PathBuf {
302        self.dir.path().join("nats-server.log")
303    }
304
305    fn signal(&self, signal: Signal) {
306        kill_process(self.pid, signal)
307            .unwrap_or_else(|err| panic!("send {signal:?} to nats-server: {err}"));
308    }
309
310    async fn wait_ready(&mut self, jetstream: bool) {
311        let deadline = Instant::now() + STARTUP_TIMEOUT;
312        let ports_path = self
313            .dir
314            .path()
315            .join(format!("nats-server_{}.ports", self.pid.as_raw_nonzero()));
316
317        // The ports file is written once all listeners are bound. It may be
318        // observed half-written, so retry until it parses.
319        let ports = loop {
320            if let Some(ports) = fs::read_to_string(&ports_path)
321                .await
322                .ok()
323                .and_then(|ports| serde_json::from_str::<Value>(&ports).ok())
324            {
325                break ports;
326            }
327            self.assert_starting(deadline).await;
328        };
329
330        let first_url = |kind: &str| ports[kind][0].as_str().map(str::to_owned);
331        self.client_url = first_url("nats").expect("ports file without a nats listener");
332        self.websocket_url = first_url("websocket");
333        self.monitoring_addr = first_url("monitoring")
334            .and_then(|url| url.strip_prefix("http://")?.parse().ok())
335            .expect("ports file without a monitoring listener");
336
337        let healthz = if jetstream {
338            "/healthz?js-enabled-only=true"
339        } else {
340            "/healthz"
341        };
342        while !matches!(
343            self.http.get(self.monitoring_addr, healthz).await,
344            Ok((StatusCode::OK, _))
345        ) {
346            self.assert_starting(deadline).await;
347        }
348    }
349
350    async fn assert_starting(&mut self, deadline: Instant) {
351        if let Some(status) = self.child.try_wait().expect("poll nats-server") {
352            panic!("nats-server exited during startup with {status}");
353        }
354        assert!(
355            Instant::now() < deadline,
356            "nats-server did not become ready within {STARTUP_TIMEOUT:?}"
357        );
358        sleep(Duration::from_millis(10)).await;
359    }
360}
361
362impl Drop for NatsServer {
363    fn drop(&mut self) {
364        // The process is reaped in the background by tokio
365        let _ = self.child.start_kill();
366
367        if thread::panicking() {
368            self.dir.disable_cleanup(true);
369            let log = std::fs::read_to_string(self.log_path()).unwrap_or_default();
370            let lines = log.lines().collect::<Vec<_>>();
371            let tail = &lines[lines.len().saturating_sub(LOG_TAIL_LINES)..];
372            eprintln!(
373                "---- nats-server log tail (full log and config kept in {}) ----\n{}",
374                self.dir.path().display(),
375                tail.join("\n")
376            );
377        }
378    }
379}
380
381fn find_nats_server() -> Option<PathBuf> {
382    if let Some(bin) = env::var_os("NATS_SERVER_BIN")
383        .map(PathBuf::from)
384        .filter(|bin| bin.is_file())
385    {
386        return Some(bin);
387    }
388
389    env::split_paths(&env::var_os("PATH")?)
390        .map(|dir| dir.join("nats-server"))
391        .find(|bin| bin.is_file())
392}
393
394/// Quote a value for the server config file, which accepts JSON strings
395fn quote(value: impl AsRef<Path>) -> String {
396    let value = value.as_ref().to_str().expect("non UTF-8 path");
397    serde_json::to_string(value).expect("serialize string")
398}
watermelon_testkit/server.rs

watermelon_testkit/
server.rs
