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                &created,
200                Some(downloads),
201            )
202            .await
203        } else {
204            test_add_crate(self.db, name, self.owner, &version, &created).await
205        }
206    }
207}
208
209/// Add a cached crate with a specified download count.
210pub async fn test_add_cached_crate_with_downloads(
211    db: &Database,
212    name: &str,
213    version: &str,
214    downloads: u64,
215) -> DbResult<()> {
216    let _ = test_add_cached_crate(db, name, version).await?;
217
218    let krate = cratesio_crate::Entity::find()
219        .filter(cratesio_crate::Column::Name.eq(name))
220        .one(&db.db_con)
221        .await?
222        .ok_or_else(|| DbError::CrateNotFound(name.to_string()))?;
223
224    let total_downloads = krate.total_downloads as u64;
225
226    let mut krate: cratesio_crate::ActiveModel = krate.into();
227    krate.total_downloads = Set((total_downloads + downloads) as i64);
228    krate.update(&db.db_con).await?;
229
230    Ok(())
231}
232
233/// Add a cached crate for testing.
234pub async fn test_add_cached_crate(db: &Database, name: &str, version: &str) -> DbResult<Prefetch> {
235    let etag = "etag";
236    let last_modified = "last_modified";
237    let description = Some(String::from("description"));
238    let indices = vec![IndexMetadata {
239        name: name.to_string(),
240        vers: version.to_string(),
241        deps: vec![],
242        cksum: "cksum".to_string(),
243        features: BTreeMap::new(),
244        features2: None,
245        pubtime: None,
246        yanked: false,
247        links: None,
248        v: Some(1),
249    }];
250
251    db.add_cratesio_prefetch_data(
252        &OriginalName::from_unchecked(name.to_string()),
253        etag,
254        last_modified,
255        description,
256        &indices,
257    )
258    .await
259}
260
261/// Add a test crate.
262pub async fn test_add_crate(
263    db: &Database,
264    name: &str,
265    owner: &str,
266    version: &Version,
267    created: &DateTime<Utc>,
268) -> DbResult<i64> {
269    let pm = PublishMetadata {
270        name: name.to_string(),
271        vers: version.to_string(),
272        ..PublishMetadata::default()
273    };
274    let user = user::Entity::find()
275        .filter(user::Column::Name.eq(owner))
276        .one(&db.db_con)
277        .await?;
278    if user.is_none() {
279        db.add_user(name, "pwd", "salt", false, false).await?;
280    }
281
282    db.add_crate(&pm, "cksum", created, owner).await
283}
284
285/// Add a test crate with a specified download count.
286pub async fn test_add_crate_with_downloads(
287    db: &Database,
288    name: &str,
289    owner: &str,
290    version: &Version,
291    created: &DateTime<Utc>,
292    downloads: Option<i64>,
293) -> DbResult<i64> {
294    let pm = PublishMetadata {
295        name: name.to_string(),
296        vers: version.to_string(),
297        ..PublishMetadata::default()
298    };
299    let user = user::Entity::find()
300        .filter(user::Column::Name.eq(owner))
301        .one(&db.db_con)
302        .await?;
303    if user.is_none() {
304        db.add_user(name, "pwd", "salt", false, false).await?;
305    }
306
307    db.add_crate(&pm, "cksum", created, owner).await?;
308    let (cm, krate) = crate_meta::Entity::find()
309        .find_also_related(krate::Entity)
310        .filter(krate::Column::Name.eq(name))
311        .filter(crate_meta::Column::Version.eq(version))
312        .one(&db.db_con)
313        .await?
314        .ok_or_else(|| DbError::CrateNotFound(name.to_string()))?;
315    let mut cm: crate_meta::ActiveModel = cm.into();
316
317    let current_downloads = krate.as_ref().unwrap().total_downloads;
318    let crate_id = krate.as_ref().unwrap().id;
319
320    let mut krate: krate::ActiveModel = krate.unwrap().into();
321    krate.total_downloads = Set(current_downloads + downloads.unwrap_or(0));
322    krate.update(&db.db_con).await?;
323    cm.downloads = Set(downloads.unwrap_or_default());
324    cm.update(&db.db_con).await?;
325    Ok(crate_id)
326}
327
328/// Add test crate metadata.
329pub async fn test_add_crate_meta(
330    db: &Database,
331    crate_id: i64,
332    version: &str,
333    created: &DateTime<Utc>,
334    downloads: Option<i64>,
335) -> DbResult<()> {
336    let cm = crate_meta::ActiveModel {
337        id: ActiveValue::default(),
338        version: Set(version.to_string()),
339        created: Set(created.to_string()),
340        downloads: Set(downloads.unwrap_or_default()),
341        crate_fk: Set(crate_id),
342        ..Default::default()
343    };
344
345    cm.insert(&db.db_con).await?;
346
347    Ok(())
348}
349
350/// Delete crate index entries for testing.
351pub async fn test_delete_crate_index(db: &Database, crate_id: i64) -> DbResult<()> {
352    crate_index::Entity::delete_many()
353        .filter(crate_index::Column::CrateFk.eq(crate_id))
354        .exec(&db.db_con)
355        .await?;
356    Ok(())
357}
358
359/// Clean database by removing old sessions.
360pub async fn clean_db(db: &Database, session_age: std::time::Duration) -> DbResult<()> {
361    let session_age = chrono::Duration::from_std(session_age).unwrap();
362    let now = std::ops::Add::add(Utc::now(), session_age)
363        .format(DB_DATE_FORMAT)
364        .to_string();
365
366    session::Entity::delete_many()
367        .filter(Expr::col(session::Column::Created).lt(now))
368        .exec(&db.db_con)
369        .await?;
370
371    Ok(())
372}
373
374/// Get crate meta list by crate ID.
375pub async fn get_crate_meta_list(db: &Database, crate_id: i64) -> DbResult<Vec<CrateMeta>> {
376    let cm: Vec<(crate_meta::Model, Option<krate::Model>)> = crate_meta::Entity::find()
377        .find_also_related(krate::Entity)
378        .filter(crate_meta::Column::CrateFk.eq(crate_id))
379        .all(&db.db_con)
380        .await?;
381
382    let crate_metas: Vec<CrateMeta> = cm
383        .into_iter()
384        .map(|(m, c)| CrateMeta {
385            name: c.unwrap().name, // Unwarp is ok, as a relation always exists
386            id: m.id,
387            version: m.version,
388            created: m.created,
389            downloads: m.downloads,
390            crate_fk: m.crate_fk,
391        })
392        .collect();
393
394    Ok(crate_metas)
395}
kellnr_db/database/test_utils.rs

kellnr_db/database/
test_utils.rs
