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

kellnr_db/database/
test_utils.rs
