kellnr_db/database/
test_utils.rs

1//! Test utilities for database testing.
2//!
3//! These functions provide convenient ways to set up test data
4//! in the database for integration and unit tests.
5//!
6//! # Builder Pattern
7//!
8//! For ergonomic test setup, use [`TestCrateBuilder`]:
9//!
10//! ```ignore
11//! TestCrateBuilder::new(test_db)
12//!     .name("mycrate")
13//!     .owner("admin")
14//!     .version("1.0.0")
15//!     .build()
16//!     .await
17//!     .unwrap();
18//! ```
19
20use std::collections::BTreeMap;
21
22use chrono::{DateTime, TimeZone, Utc};
23use kellnr_common::index_metadata::IndexMetadata;
24use kellnr_common::original_name::OriginalName;
25use kellnr_common::prefetch::Prefetch;
26use kellnr_common::publish_metadata::PublishMetadata;
27use kellnr_common::version::Version;
28use kellnr_entity::{
29    auth_token, crate_index, crate_meta, cratesio_crate, krate, oauth2_state, session, user,
30};
31use sea_orm::sea_query::Expr;
32use sea_orm::{
33    ActiveModelTrait, ActiveValue, ColumnTrait, EntityTrait, ExprTrait, QueryFilter, Set,
34};
35
36use super::{DB_DATE_FORMAT, Database, parse_db_version};
37use crate::CrateMeta;
38use crate::error::DbError;
39use crate::provider::{DbProvider, DbResult};
40
41/// Returns a standard test date for consistent test data.
42///
43/// The returned date is `2020-10-07 13:18:00 UTC`, which is commonly
44/// used across the test suite.
45#[must_use]
46pub fn default_created() -> DateTime<Utc> {
47    Utc.with_ymd_and_hms(2020, 10, 7, 13, 18, 0).unwrap()
48}
49
50/// Add multiple versions of a crate at once.
51///
52/// This is a convenience function for tests that need to set up a crate
53/// with multiple versions quickly.
54///
55/// # Arguments
56///
57/// * `db` - The database connection
58/// * `name` - The crate name
59/// * `owner` - The owner username
60/// * `versions` - Slice of version strings to add (e.g., `["1.0.0", "2.0.0"]`)
61///
62/// # Returns
63///
64/// A vector of crate IDs for each version added.
65///
66/// # Example
67///
68/// ```ignore
69/// let ids = add_multiple_versions(test_db, "mycrate", "admin", &["1.0.0", "2.0.0", "3.0.0"])
70///     .await
71///     .unwrap();
72/// ```
73pub async fn add_multiple_versions(
74    db: &Database,
75    name: &str,
76    owner: &str,
77    versions: &[&str],
78) -> DbResult<Vec<i64>> {
79    let created = default_created();
80    let mut ids = Vec::with_capacity(versions.len());
81    for version in versions {
82        let id = test_add_crate(db, name, owner, &parse_db_version(version)?, &created).await?;
83        ids.push(id);
84    }
85    Ok(ids)
86}
87
88/// A builder for creating test crates with a fluent API.
89///
90/// # Example
91///
92/// ```ignore
93/// // Minimal usage with defaults
94/// TestCrateBuilder::new(test_db)
95///     .name("mycrate")
96///     .build()
97///     .await
98///     .unwrap();
99///
100/// // Full customization
101/// let crate_id = TestCrateBuilder::new(test_db)
102///     .name("mycrate")
103///     .owner("testuser")
104///     .version("2.0.0")
105///     .created(Utc::now())
106///     .downloads(100)
107///     .build()
108///     .await
109///     .unwrap();
110/// ```
111pub struct TestCrateBuilder<'a> {
112    db: &'a Database,
113    name: Option<&'a str>,
114    owner: &'a str,
115    version: &'a str,
116    created: Option<DateTime<Utc>>,
117    downloads: Option<i64>,
118}
119
120impl<'a> TestCrateBuilder<'a> {
121    /// Create a new builder with default values.
122    ///
123    /// Defaults:
124    /// - `owner`: "admin"
125    /// - `version`: "1.0.0"
126    /// - `created`: [`default_created()`]
127    /// - `downloads`: None (0)
128    #[must_use]
129    pub fn new(db: &'a Database) -> Self {
130        Self {
131            db,
132            name: None,
133            owner: "admin",
134            version: "1.0.0",
135            created: None,
136            downloads: None,
137        }
138    }
139
140    /// Set the crate name (required).
141    #[must_use]
142    pub fn name(mut self, name: &'a str) -> Self {
143        self.name = Some(name);
144        self
145    }
146
147    /// Set the owner username. Defaults to "admin".
148    #[must_use]
149    pub fn owner(mut self, owner: &'a str) -> Self {
150        self.owner = owner;
151        self
152    }
153
154    /// Set the version string. Defaults to "1.0.0".
155    #[must_use]
156    pub fn version(mut self, version: &'a str) -> Self {
157        self.version = version;
158        self
159    }
160
161    /// Set the creation date. Defaults to [`default_created()`].
162    #[must_use]
163    pub fn created(mut self, created: DateTime<Utc>) -> Self {
164        self.created = Some(created);
165        self
166    }
167
168    /// Set the download count. If not set, downloads will be 0.
169    #[must_use]
170    pub fn downloads(mut self, downloads: i64) -> Self {
171        self.downloads = Some(downloads);
172        self
173    }
174
175    /// Build and insert the crate into the database.
176    ///
177    /// # Errors
178    ///
179    /// Returns an error if:
180    /// - The crate name was not set
181    /// - The version string is invalid
182    /// - Database insertion fails
183    ///
184    /// # Returns
185    ///
186    /// The crate ID on success.
187    pub async fn build(self) -> DbResult<i64> {
188        let name = self.name.ok_or_else(|| {
189            DbError::InvalidCrateName("TestCrateBuilder: name is required".to_string())
190        })?;
191        let created = self.created.unwrap_or_else(default_created);
192        let version = Version::try_from(self.version)
193            .map_err(|_| DbError::InvalidVersion(self.version.to_string()))?;
194
195        if let Some(downloads) = self.downloads {
196            test_add_crate_with_downloads(
197                self.db,
198                name,
199                self.owner,
200                &version,
201                None,
202                &created,
203                Some(downloads),
204            )
205            .await
206        } else {
207            test_add_crate(self.db, name, self.owner, &version, &created).await
208        }
209    }
210}
211
212/// Add a cached crate with a specified download count.
213pub async fn test_add_cached_crate_with_downloads(
214    db: &Database,
215    name: &str,
216    version: &str,
217    downloads: u64,
218) -> DbResult<()> {
219    let _ = test_add_cached_crate(db, name, version).await?;
220
221    let krate = cratesio_crate::Entity::find()
222        .filter(cratesio_crate::Column::Name.eq(name))
223        .one(&db.db_con)
224        .await?
225        .ok_or_else(|| DbError::CrateNotFound(name.to_string()))?;
226
227    let total_downloads = krate.total_downloads as u64;
228
229    let mut krate: cratesio_crate::ActiveModel = krate.into();
230    krate.total_downloads = Set((total_downloads + downloads) as i64);
231    krate.update(&db.db_con).await?;
232
233    Ok(())
234}
235
236/// Add a cached crate for testing.
237pub async fn test_add_cached_crate(db: &Database, name: &str, version: &str) -> DbResult<Prefetch> {
238    let etag = "etag";
239    let last_modified = "last_modified";
240    let description = Some(String::from("description"));
241    let indices = vec![IndexMetadata {
242        name: name.to_string(),
243        vers: version.to_string(),
244        deps: vec![],
245        cksum: "cksum".to_string(),
246        features: BTreeMap::new(),
247        features2: None,
248        pubtime: None,
249        yanked: false,
250        links: None,
251        v: Some(1),
252        rust_version: None,
253    }];
254
255    db.add_cratesio_prefetch_data(
256        &OriginalName::from_unchecked(name.to_string()),
257        etag,
258        last_modified,
259        description,
260        &indices,
261    )
262    .await
263}
264
265/// Add a test crate.
266pub async fn test_add_crate(
267    db: &Database,
268    name: &str,
269    owner: &str,
270    version: &Version,
271    created: &DateTime<Utc>,
272) -> DbResult<i64> {
273    let pm = PublishMetadata {
274        name: name.to_string(),
275        vers: version.to_string(),
276        ..PublishMetadata::default()
277    };
278    let user = user::Entity::find()
279        .filter(user::Column::Name.eq(owner))
280        .one(&db.db_con)
281        .await?;
282    if user.is_none() {
283        db.add_user(name, "pwd", false, false).await?;
284    }
285
286    db.add_crate(&pm, "cksum", 1024, created, owner, None).await
287}
288
289/// Insert an empty crate without any name checks, to simulate data created
290/// before a check was introduced.
291pub async fn test_insert_crate_unchecked(db: &Database, name: &str) -> DbResult<i64> {
292    let krate = krate::ActiveModel {
293        id: ActiveValue::default(),
294        name: Set(name.to_lowercase()),
295        original_name: Set(name.to_string()),
296        max_version: Set("0.0.0".to_string()),
297        last_updated: Set(default_created().format(DB_DATE_FORMAT).to_string()),
298        total_downloads: Set(0),
299        homepage: Set(None),
300        description: Set(None),
301        repository: Set(None),
302        e_tag: Set(String::new()),
303        restricted_download: Set(false),
304    };
305    Ok(krate.insert(&db.db_con).await?.id)
306}
307
308/// Add a test crate with a specified download count.
309pub async fn test_add_crate_with_downloads(
310    db: &Database,
311    name: &str,
312    owner: &str,
313    version: &Version,
314    description: Option<&str>,
315    created: &DateTime<Utc>,
316    downloads: Option<i64>,
317) -> DbResult<i64> {
318    let pm = PublishMetadata {
319        name: name.to_string(),
320        vers: version.to_string(),
321        description: description.map(ToString::to_string),
322        ..PublishMetadata::default()
323    };
324    let user = user::Entity::find()
325        .filter(user::Column::Name.eq(owner))
326        .one(&db.db_con)
327        .await?;
328    if user.is_none() {
329        db.add_user(name, "pwd", false, false).await?;
330    }
331
332    db.add_crate(&pm, "cksum", 1024, created, owner, None)
333        .await?;
334    let (cm, krate) = crate_meta::Entity::find()
335        .find_also_related(krate::Entity)
336        .filter(krate::Column::Name.eq(name))
337        .filter(crate_meta::Column::Version.eq(version))
338        .one(&db.db_con)
339        .await?
340        .ok_or_else(|| DbError::CrateNotFound(name.to_string()))?;
341    let mut cm: crate_meta::ActiveModel = cm.into();
342
343    let current_downloads = krate.as_ref().unwrap().total_downloads;
344    let crate_id = krate.as_ref().unwrap().id;
345
346    let mut krate: krate::ActiveModel = krate.unwrap().into();
347    krate.total_downloads = Set(current_downloads + downloads.unwrap_or(0));
348    krate.update(&db.db_con).await?;
349    cm.downloads = Set(downloads.unwrap_or_default());
350    cm.update(&db.db_con).await?;
351    Ok(crate_id)
352}
353
354/// Add test crate metadata.
355pub async fn test_add_crate_meta(
356    db: &Database,
357    crate_id: i64,
358    version: &str,
359    created: &DateTime<Utc>,
360    downloads: Option<i64>,
361) -> DbResult<()> {
362    let cm = crate_meta::ActiveModel {
363        id: ActiveValue::default(),
364        version: Set(version.to_string()),
365        created: Set(created.to_string()),
366        downloads: Set(downloads.unwrap_or_default()),
367        crate_fk: Set(crate_id),
368        ..Default::default()
369    };
370
371    cm.insert(&db.db_con).await?;
372
373    Ok(())
374}
375
376/// Delete crate index entries for testing.
377pub async fn test_delete_crate_index(db: &Database, crate_id: i64) -> DbResult<()> {
378    crate_index::Entity::delete_many()
379        .filter(crate_index::Column::CrateFk.eq(crate_id))
380        .exec(&db.db_con)
381        .await?;
382    Ok(())
383}
384
385/// Clean database by removing old sessions.
386pub async fn clean_db(db: &Database, session_age: std::time::Duration) -> DbResult<()> {
387    let session_age = chrono::Duration::from_std(session_age).unwrap();
388    let now = std::ops::Add::add(Utc::now(), session_age)
389        .format(DB_DATE_FORMAT)
390        .to_string();
391
392    session::Entity::delete_many()
393        .filter(Expr::col(session::Column::Created).lt(now))
394        .exec(&db.db_con)
395        .await?;
396
397    Ok(())
398}
399
400/// Insert a session whose `created` timestamp is `age` in the past.
401///
402/// Used to test session-expiry behavior without sleeping: pass an `age`
403/// larger than the configured session age to create an already-expired
404/// session, or a small one to create a still-valid session.
405pub async fn add_aged_session(
406    db: &Database,
407    user_name: &str,
408    token: &str,
409    age: std::time::Duration,
410) -> DbResult<()> {
411    let user = user::Entity::find()
412        .filter(user::Column::Name.eq(user_name))
413        .one(&db.db_con)
414        .await?
415        .ok_or_else(|| DbError::UserNotFound(user_name.to_owned()))?;
416
417    let age = chrono::Duration::from_std(age).expect("session age out of range");
418    let created = (Utc::now() - age).format(DB_DATE_FORMAT).to_string();
419    let s = session::ActiveModel {
420        token: Set(crate::password::hash_token(token)),
421        created: Set(created),
422        user_fk: Set(user.id),
423        ..Default::default()
424    };
425    s.insert(&db.db_con).await?;
426    Ok(())
427}
428
429/// Get crate meta list by crate ID.
430pub async fn get_crate_meta_list(db: &Database, crate_id: i64) -> DbResult<Vec<CrateMeta>> {
431    let cm: Vec<(crate_meta::Model, Option<krate::Model>)> = crate_meta::Entity::find()
432        .find_also_related(krate::Entity)
433        .filter(crate_meta::Column::CrateFk.eq(crate_id))
434        .all(&db.db_con)
435        .await?;
436
437    let crate_metas: Vec<CrateMeta> = cm
438        .into_iter()
439        .map(|(m, c)| CrateMeta {
440            name: c.unwrap().name, // Unwarp is ok, as a relation always exists
441            id: m.id,
442            version: m.version,
443            created: m.created,
444            downloads: m.downloads,
445            crate_fk: m.crate_fk,
446        })
447        .collect();
448
449    Ok(crate_metas)
450}
451
452/// Stores a legacy `sha256(pwd + salt)` password hash for `user_name`, as
453/// written by Kellnr versions before the switch to argon2.
454pub async fn set_legacy_pwd(db: &Database, user_name: &str, pwd: &str, salt: &str) -> DbResult<()> {
455    let mut u: user::ActiveModel = db.get_user_model(user_name).await?.into();
456    u.pwd = Set(crate::password::legacy_hash(pwd, salt));
457    u.salt = Set(salt.to_owned());
458    u.update(&db.db_con).await?;
459    Ok(())
460}
461
462/// Sets the last use of every auth token named `name`, e.g. to a date in the past.
463pub async fn set_auth_token_last_used(db: &Database, name: &str, last_used: &str) -> DbResult<()> {
464    auth_token::Entity::update_many()
465        .col_expr(auth_token::Column::LastUsed, Expr::value(last_used))
466        .filter(auth_token::Column::Name.eq(name))
467        .exec(&db.db_con)
468        .await?;
469    Ok(())
470}
471
472/// Returns the session tokens exactly as stored in the database.
473pub async fn stored_session_tokens(db: &Database) -> DbResult<Vec<String>> {
474    Ok(session::Entity::find()
475        .all(&db.db_con)
476        .await?
477        .into_iter()
478        .map(|s| s.token)
479        .collect())
480}
481
482/// Stores an `OAuth2` state whose `created` timestamp lies `age` in the past.
483pub async fn add_aged_oauth2_state(
484    db: &Database,
485    state: &str,
486    age: std::time::Duration,
487) -> DbResult<()> {
488    let age = chrono::Duration::from_std(age).expect("state age out of range");
489    let s = oauth2_state::ActiveModel {
490        state: Set(state.to_owned()),
491        pkce_verifier: Set("verifier".to_owned()),
492        nonce: Set("nonce".to_owned()),
493        created: Set((Utc::now() - age).format(DB_DATE_FORMAT).to_string()),
494        ..Default::default()
495    };
496    s.insert(&db.db_con).await?;
497    Ok(())
498}
kellnr_db/database/test_utils.rs

kellnr_db/database/
test_utils.rs
