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, 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", "salt", false, false).await?;
282    }
283
284    db.add_crate(&pm, "cksum", created, owner).await
285}
286
287/// Add a test crate with a specified download count.
288pub async fn test_add_crate_with_downloads(
289    db: &Database,
290    name: &str,
291    owner: &str,
292    version: &Version,
293    description: Option<&str>,
294    created: &DateTime<Utc>,
295    downloads: Option<i64>,
296) -> DbResult<i64> {
297    let pm = PublishMetadata {
298        name: name.to_string(),
299        vers: version.to_string(),
300        description: description.map(ToString::to_string),
301        ..PublishMetadata::default()
302    };
303    let user = user::Entity::find()
304        .filter(user::Column::Name.eq(owner))
305        .one(&db.db_con)
306        .await?;
307    if user.is_none() {
308        db.add_user(name, "pwd", "salt", false, false).await?;
309    }
310
311    db.add_crate(&pm, "cksum", created, owner).await?;
312    let (cm, krate) = crate_meta::Entity::find()
313        .find_also_related(krate::Entity)
314        .filter(krate::Column::Name.eq(name))
315        .filter(crate_meta::Column::Version.eq(version))
316        .one(&db.db_con)
317        .await?
318        .ok_or_else(|| DbError::CrateNotFound(name.to_string()))?;
319    let mut cm: crate_meta::ActiveModel = cm.into();
320
321    let current_downloads = krate.as_ref().unwrap().total_downloads;
322    let crate_id = krate.as_ref().unwrap().id;
323
324    let mut krate: krate::ActiveModel = krate.unwrap().into();
325    krate.total_downloads = Set(current_downloads + downloads.unwrap_or(0));
326    krate.update(&db.db_con).await?;
327    cm.downloads = Set(downloads.unwrap_or_default());
328    cm.update(&db.db_con).await?;
329    Ok(crate_id)
330}
331
332/// Add test crate metadata.
333pub async fn test_add_crate_meta(
334    db: &Database,
335    crate_id: i64,
336    version: &str,
337    created: &DateTime<Utc>,
338    downloads: Option<i64>,
339) -> DbResult<()> {
340    let cm = crate_meta::ActiveModel {
341        id: ActiveValue::default(),
342        version: Set(version.to_string()),
343        created: Set(created.to_string()),
344        downloads: Set(downloads.unwrap_or_default()),
345        crate_fk: Set(crate_id),
346        ..Default::default()
347    };
348
349    cm.insert(&db.db_con).await?;
350
351    Ok(())
352}
353
354/// Delete crate index entries for testing.
355pub async fn test_delete_crate_index(db: &Database, crate_id: i64) -> DbResult<()> {
356    crate_index::Entity::delete_many()
357        .filter(crate_index::Column::CrateFk.eq(crate_id))
358        .exec(&db.db_con)
359        .await?;
360    Ok(())
361}
362
363/// Clean database by removing old sessions.
364pub async fn clean_db(db: &Database, session_age: std::time::Duration) -> DbResult<()> {
365    let session_age = chrono::Duration::from_std(session_age).unwrap();
366    let now = std::ops::Add::add(Utc::now(), session_age)
367        .format(DB_DATE_FORMAT)
368        .to_string();
369
370    session::Entity::delete_many()
371        .filter(Expr::col(session::Column::Created).lt(now))
372        .exec(&db.db_con)
373        .await?;
374
375    Ok(())
376}
377
378/// Insert a session whose `created` timestamp is `age` in the past.
379///
380/// Used to test session-expiry behavior without sleeping: pass an `age`
381/// larger than the configured session age to create an already-expired
382/// session, or a small one to create a still-valid session.
383pub async fn add_aged_session(
384    db: &Database,
385    user_name: &str,
386    token: &str,
387    age: std::time::Duration,
388) -> DbResult<()> {
389    let user = user::Entity::find()
390        .filter(user::Column::Name.eq(user_name))
391        .one(&db.db_con)
392        .await?
393        .ok_or_else(|| DbError::UserNotFound(user_name.to_owned()))?;
394
395    let age = chrono::Duration::from_std(age).expect("session age out of range");
396    let created = (Utc::now() - age).format(DB_DATE_FORMAT).to_string();
397    let s = session::ActiveModel {
398        token: Set(token.to_owned()),
399        created: Set(created),
400        user_fk: Set(user.id),
401        ..Default::default()
402    };
403    s.insert(&db.db_con).await?;
404    Ok(())
405}
406
407/// Get crate meta list by crate ID.
408pub async fn get_crate_meta_list(db: &Database, crate_id: i64) -> DbResult<Vec<CrateMeta>> {
409    let cm: Vec<(crate_meta::Model, Option<krate::Model>)> = crate_meta::Entity::find()
410        .find_also_related(krate::Entity)
411        .filter(crate_meta::Column::CrateFk.eq(crate_id))
412        .all(&db.db_con)
413        .await?;
414
415    let crate_metas: Vec<CrateMeta> = cm
416        .into_iter()
417        .map(|(m, c)| CrateMeta {
418            name: c.unwrap().name, // Unwarp is ok, as a relation always exists
419            id: m.id,
420            version: m.version,
421            created: m.created,
422            downloads: m.downloads,
423            crate_fk: m.crate_fk,
424        })
425        .collect();
426
427    Ok(crate_metas)
428}
kellnr_db/database/test_utils.rs

kellnr_db/database/
test_utils.rs
