Skip to content

Commit

Permalink
crypto: Allow duplicating a Client in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
andybalaam committed May 13, 2024
1 parent 2c57557 commit 749ed2c
Showing 1 changed file with 64 additions and 7 deletions.
71 changes: 64 additions & 7 deletions testing/matrix-sdk-integration-testing/src/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::{
collections::HashMap,
ops::Deref,
option_env,
path::{Path, PathBuf},
sync::{Arc, Mutex as StdMutex},
time::Duration,
};
Expand All @@ -21,9 +22,14 @@ use tokio::sync::Mutex;

static USERS: Lazy<Mutex<HashMap<String, (Client, TempDir)>>> = Lazy::new(Mutex::default);

enum SqlitePath {
Random,
Path(PathBuf),
}

pub struct TestClientBuilder {
username: String,
use_sqlite: bool,
use_sqlite_dir: Option<SqlitePath>,
encryption_settings: EncryptionSettings,
http_proxy: Option<String>,
}
Expand All @@ -32,7 +38,7 @@ impl TestClientBuilder {
pub fn new(username: impl Into<String>) -> Self {
Self {
username: username.into(),
use_sqlite: false,
use_sqlite_dir: None,
encryption_settings: Default::default(),
http_proxy: None,
}
Expand All @@ -45,7 +51,16 @@ impl TestClientBuilder {
}

pub fn use_sqlite(mut self) -> Self {
self.use_sqlite = true;
self.use_sqlite_dir = Some(SqlitePath::Random);
self
}

/// Create or re-use a Sqlite store (with no passphrase) in the supplied
/// directory. Note: this path must remain valid throughout the use of
/// the constructed Client, so if you created a TempDir you must hang on
/// to a reference to it throughout the test.
pub fn use_sqlite_dir(mut self, path: &Path) -> Self {
self.use_sqlite_dir = Some(SqlitePath::Path(path.to_owned()));
self
}

Expand All @@ -59,6 +74,44 @@ impl TestClientBuilder {
self
}

/// Create a new Client that is a copy of the supplied one, created using
/// [`Client::restore_session`].
pub async fn duplicate(self, other: &Client) -> Result<Client> {
let homeserver_url =
option_env!("HOMESERVER_URL").unwrap_or("http://localhost:8228").to_owned();
let sliding_sync_proxy_url =
option_env!("SLIDING_SYNC_PROXY_URL").unwrap_or("http://localhost:8338").to_owned();

let mut client_builder = Client::builder()
.user_agent("matrix-sdk-integration-tests")
.homeserver_url(homeserver_url)
.sliding_sync_proxy(sliding_sync_proxy_url)
.with_encryption_settings(self.encryption_settings)
.request_config(RequestConfig::short_retry());

if let Some(proxy) = self.http_proxy {
client_builder = client_builder.proxy(proxy);
}

let client = match self.use_sqlite_dir {
Some(SqlitePath::Path(path_buf)) => {
client_builder.sqlite_store(&path_buf, None).build().await?
}
_ => {
panic!("You must call use_sqlite_dir for a duplicate client!");
}
};

client
.restore_session(
other.session().expect("Session must be logged in before we can duplicate it"),
)
.await
.expect("Failed to restore session");

Ok(client)
}

pub async fn build(self) -> Result<Client> {
let mut users = USERS.lock().await;
if let Some((client, _)) = users.get(&self.username) {
Expand All @@ -83,10 +136,14 @@ impl TestClientBuilder {
client_builder = client_builder.proxy(proxy);
}

let client = if self.use_sqlite {
client_builder.sqlite_store(tmp_dir.path(), None).build().await?
} else {
client_builder.build().await?
let client = match self.use_sqlite_dir {
None => client_builder.build().await?,
Some(SqlitePath::Random) => {
client_builder.sqlite_store(tmp_dir.path(), None).build().await?
}
Some(SqlitePath::Path(path_buf)) => {
client_builder.sqlite_store(&path_buf, None).build().await?
}
};

// safe to assume we have not registered this user yet, but ignore if we did
Expand Down

0 comments on commit 749ed2c

Please sign in to comment.