CI - rust smoketests lib expansion

crates/smoketests/src/lib.rs

@@ -58,6 +58,7 @@ use regex::Regex;
 use spacetimedb_guard::{ensure_binaries_built, SpacetimeDbGuard};
 use std::env;
 use std::fs;
+use std::io::{BufRead, BufReader};
 use std::path::{Path, PathBuf};
 use std::process::{Command, Output, Stdio};
 use std::sync::OnceLock;
@@ -401,6 +402,16 @@ impl ApiResponse {
     }
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct PublishOptions {
+    pub clear: bool,
+    pub break_clients: bool,
+    pub num_replicas: Option<u32>,
+    pub organization: Option<String>,
+    pub force: bool,
+    pub stdin_input: Option<&str>,
+}
+
 /// Builder for creating `Smoketest` instances.
 pub struct SmoketestBuilder {
     module_code: Option<String>,
@@ -409,6 +420,7 @@ pub struct SmoketestBuilder {
     extra_deps: String,
     autopublish: bool,
     pg_port: Option<u16>,
+    server_url_override: Option<String>,
 }
 
 impl Default for SmoketestBuilder {
@@ -427,9 +439,15 @@ impl SmoketestBuilder {
             extra_deps: String::new(),
             autopublish: true,
             pg_port: None,
+            server_url_override: None,
         }
     }
 
+    pub fn server_url(mut self, url: &str) -> Self {
+        self.server_url_override = Some(url.to_string());
+        self
+    }
+
     /// Enables the PostgreSQL wire protocol on the specified port.
     pub fn pg_port(mut self, port: u16) -> Self {
         self.pg_port = Some(port);
@@ -503,7 +521,10 @@ impl SmoketestBuilder {
         let build_start = Instant::now();
 
         // Check if we're running against a remote server
-        let (guard, server_url) = if let Some(remote_url) = remote_server_url() {
+        let (guard, server_url) = if let Some(url) = self.server_url_override {
+            eprintln!("[REMOTE] Using explicit server URL: {}", url);
+            (None, url)
+        } else if let Some(remote_url) = remote_server_url() {
             eprintln!("[REMOTE] Using remote server: {}", remote_url);
             (None, remote_url)
         } else {
@@ -629,6 +650,16 @@ impl Smoketest {
             .context("No spacetimedb_token found in config")
     }
 
+    pub fn login_with_token(&self, token: &str) -> Result<()> {
+        let host = self.server_host();
+        let config_str = format!(
+            "default_server = \"localhost\"\n\nspacetimedb_token = \"{}\"\n\n[[server_configs]]\nnickname = \"localhost\"\nhost = \"{}\"\nprotocol = \"http\"\n",
+            token, host
+        );
+        fs::write(&self.config_path, config_str).context("Failed to write config.toml")?;
+        Ok(())
+    }
+
     /// Runs psql command against the PostgreSQL wire protocol server.
     ///
     /// Returns the output on success, or an error with stderr on failure.
@@ -1009,6 +1040,10 @@ log = "0.4"
         self.publish_module_opts(Some(name), clear)
     }
 
+    pub fn publish_module_named_ext(&mut self, name: &str, opts: PublishOptions) -> Result<String> {
+        self.publish_module_internal_ext(Some(name), opts)
+    }
+
     /// Re-publishes the module to the existing database identity with optional clear.
     ///
     /// This is useful for testing auto-migrations where you want to update
@@ -1040,20 +1075,28 @@ log = "0.4"
         self.publish_module_internal(Some(name), false, false, false, None)
     }
 
+    pub fn publish_module_with_options_ext(&mut self, name: &str, opts: PublishOptions) -> Result<String> {
+        self.publish_module_internal_ext(Some(name), opts)
+    }
+
     /// Internal helper for publishing with options.
     fn publish_module_opts(&mut self, name: Option<&str>, clear: bool) -> Result<String> {
         self.publish_module_internal(name, clear, false, true, None)
     }
 
     /// Internal helper for publishing with all options.
-    fn publish_module_internal(
-        &mut self,
-        name: Option<&str>,
-        clear: bool,
-        break_clients: bool,
-        force: bool,
-        stdin_input: Option<&str>,
-    ) -> Result<String> {
+    fn publish_module_internal(&mut self, name: Option<&str>, clear: bool, break_clients: bool) -> Result<String> {
+        self.publish_module_internal_ext(
+            name,
+            PublishOptions {
+                clear,
+                break_clients,
+                ..PublishOptions::default()
+            },
+        )
+    }
+
+    fn publish_module_internal_ext(&mut self, name: Option<&str>, opts: PublishOptions) -> Result<String> {
         let start = Instant::now();
 
         // Determine the WASM path - either precompiled or build it
@@ -1100,14 +1143,28 @@ log = "0.4"
             args.push("--yes");
         }
 
-        if clear {
+        if opts.clear {
             args.push("--clear-database");
         }
 
-        if break_clients {
+        if opts.break_clients {
             args.push("--break-clients");
         }
 
+        let mut num_replicas_owned: Option<String> = None;
+        if let Some(n) = opts.num_replicas {
+            num_replicas_owned = Some(n.to_string());
+            args.push("--num-replicas");
+            args.push(num_replicas_owned.as_ref().unwrap());
+        }
+
+        let mut org_owned: Option<String> = None;
+        if let Some(org) = opts.organization.as_ref() {
+            org_owned = Some(org.clone());
+            args.push("--organization");
+            args.push(org_owned.as_ref().unwrap());
+        }
+
         let name_owned;
         if let Some(n) = name {
             name_owned = n.to_string();
@@ -1396,15 +1453,45 @@ log = "0.4"
         self.subscribe_opts(queries, n, None)
     }
 
+    pub fn subscribe_on(&self, database: &str, queries: &[&str], n: usize) -> Result<Vec<serde_json::Value>> {
+        self.subscribe_on_opts(database, queries, n, false)
+    }
+
     /// Starts a subscription with --confirmed flag and waits for N updates.
     pub fn subscribe_confirmed(&self, queries: &[&str], n: usize) -> Result<Vec<serde_json::Value>> {
         self.subscribe_opts(queries, n, Some(true))
     }
 
+    pub fn subscribe_on_confirmed(&self, database: &str, queries: &[&str], n: usize) -> Result<Vec<serde_json::Value>> {
+        self.subscribe_on_opts(database, queries, n, true)
+    }
+
     /// Internal helper for subscribe with options.
     fn subscribe_opts(&self, queries: &[&str], n: usize, confirmed: Option<bool>) -> Result<Vec<serde_json::Value>> {
         let start = Instant::now();
         let identity = self.database_identity.as_ref().context("No database published")?;
+        self.subscribe_on_impl(identity, queries, n, confirmed, start)
+    }
+
+    fn subscribe_on_opts(
+        &self,
+        database: &str,
+        queries: &[&str],
+        n: usize,
+        confirmed: Option<bool>,
+    ) -> Result<Vec<serde_json::Value>> {
+        let start = Instant::now();
+        self.subscribe_on_impl(database, queries, n, confirmed, start)
+    }
+
+    fn subscribe_on_impl(
+        &self,
+        database: &str,
+        queries: &[&str],
+        n: usize,
+        confirmed: Option<bool>,
+        start: Instant,
+    ) -> Result<Vec<serde_json::Value>> {
         let config_path_str = self.config_path.to_str().unwrap();
 
         let cli_path = ensure_binaries_built();
@@ -1415,7 +1502,7 @@ log = "0.4"
             "subscribe".to_string(),
             "--server".to_string(),
             self.server_url.to_string(),
-            identity.to_string(),
+            database.to_string(),
             "-t".to_string(),
             "30".to_string(),
             "-n".to_string(),
@@ -1456,6 +1543,10 @@ log = "0.4"
         self.subscribe_background_opts(queries, n, None)
     }
 
+    pub fn subscribe_background_on(&self, database: &str, queries: &[&str], n: usize) -> Result<SubscriptionHandle> {
+        self.subscribe_background_on_opts(database, queries, n, false)
+    }
+
     /// Starts a subscription in the background with --confirmed flag.
     pub fn subscribe_background_confirmed(&self, queries: &[&str], n: usize) -> Result<SubscriptionHandle> {
         self.subscribe_background_opts(queries, n, Some(true))
@@ -1466,21 +1557,48 @@ log = "0.4"
         self.subscribe_background_opts(queries, n, Some(false))
     }
 
+    pub fn subscribe_background_on_confirmed(
+        &self,
+        database: &str,
+        queries: &[&str],
+        n: usize,
+    ) -> Result<SubscriptionHandle> {
+        self.subscribe_background_on_opts(database, queries, n, true)
+    }
+
     /// Internal helper for background subscribe with options.
     fn subscribe_background_opts(
         &self,
         queries: &[&str],
         n: usize,
         confirmed: Option<bool>,
     ) -> Result<SubscriptionHandle> {
-        use std::io::{BufRead, BufReader};
-
         let identity = self
             .database_identity
             .as_ref()
             .context("No database published")?
             .clone();
 
+        self.subscribe_background_on_impl(&identity, queries, n, confirmed)
+    }
+
+    fn subscribe_background_on_opts(
+        &self,
+        database: &str,
+        queries: &[&str],
+        n: usize,
+        confirmed: Option<bool>,
+    ) -> Result<SubscriptionHandle> {
+        self.subscribe_background_on_impl(database, queries, n, confirmed)
+    }
+
+    fn subscribe_background_on_impl(
+        &self,
+        database: &str,
+        queries: &[&str],
+        n: usize,
+        confirmed: Option<bool>,
+    ) -> Result<SubscriptionHandle> {
         let cli_path = ensure_binaries_built();
         let mut cmd = Command::new(&cli_path);
         // Use --print-initial-update so we know when subscription is established
@@ -1491,7 +1609,7 @@ log = "0.4"
             "subscribe".to_string(),
             "--server".to_string(),
             self.server_url.clone(),
-            identity,
+            database.to_string(),
             "-t".to_string(),
             "30".to_string(),
             "-n".to_string(),

smoketests/tests/replication.py

@@ -96,11 +96,13 @@ def wait_for_leader_change(self, previous_leader_node, max_attempts=10, delay=2)
             try:
                  current_leader_node = self.get_leader_info()['node_id']
                  if current_leader_node != previous_leader_node:
+                    print(f"Leader changed from {previous_leader_node} to {current_leader_node}")
                     return current_leader_node
             except Exception:
                  print("No current leader")
 
             time.sleep(delay)
+        print('leader change timeout')
         return None
 
     def ensure_leader_health(self, id):