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

kellnr_db/database/
test_utils.rs
