use std::any::TypeId;
use std::collections::HashMap;
use std::marker::PhantomData;
use std::sync::{Mutex, OnceLock};
use crate::model::repository::insert_entity;
use crate::model::{DbError, DbPool, Model};
pub trait Factory: Model + Sized {
fn factory_default(seq: u64) -> Self;
}
pub trait FactoryValue: Sized {
fn factory_value(field: &str, seq: u64) -> Self;
fn factory_key(field: &str, seq: u64) -> Self {
Self::factory_value(field, seq)
}
}
macro_rules! numbers {
($($t:ty),*) => {$(
impl FactoryValue for $t {
fn factory_value(_field: &str, _seq: u64) -> Self {
0 as $t
}
fn factory_key(_field: &str, seq: u64) -> Self {
seq as $t
}
}
)*};
}
numbers!(i8, i16, i32, i64, u8, u16, u32, u64, f32, f64);
impl FactoryValue for bool {
fn factory_value(_field: &str, _seq: u64) -> Self {
false
}
}
impl FactoryValue for String {
fn factory_value(field: &str, seq: u64) -> Self {
if field.to_ascii_lowercase().contains("email") {
format!("{field}-{seq}@example.com")
} else {
format!("{field}-{seq}")
}
}
}
impl<T> FactoryValue for Option<T> {
fn factory_value(_field: &str, _seq: u64) -> Self {
None
}
}
impl FactoryValue for Vec<u8> {
fn factory_value(_field: &str, _seq: u64) -> Self {
Vec::new()
}
}
fn next_seq<T: 'static>() -> u64 {
static SEQUENCES: OnceLock<Mutex<HashMap<TypeId, u64>>> = OnceLock::new();
let mut sequences = SEQUENCES.get_or_init(Default::default).lock().unwrap_or_else(|e| e.into_inner());
let seq = sequences.entry(TypeId::of::<T>()).or_insert(0);
*seq += 1;
*seq
}
type Override<T> = Box<dyn Fn(&mut T) + Send + Sync>;
pub struct ModelFactory<T> {
overrides: Vec<Override<T>>,
_model: PhantomData<fn() -> T>,
}
impl<T: Factory + 'static> Default for ModelFactory<T> {
fn default() -> Self {
ModelFactory::new()
}
}
impl<T: Factory + 'static> ModelFactory<T> {
pub fn new() -> Self {
ModelFactory { overrides: Vec::new(), _model: PhantomData }
}
pub fn with(mut self, f: impl Fn(&mut T) + Send + Sync + 'static) -> Self {
self.overrides.push(Box::new(f));
self
}
pub fn build(&self) -> T {
let mut model = T::factory_default(next_seq::<T>());
for f in &self.overrides {
f(&mut model);
}
model
}
pub fn build_many(&self, n: usize) -> Vec<T> {
(0..n).map(|_| self.build()).collect()
}
pub async fn create(&self, pool: &DbPool) -> Result<T, DbError> {
insert_entity(pool, &self.build()).await
}
pub async fn create_many(&self, n: usize, pool: &DbPool) -> Result<Vec<T>, DbError> {
let mut created = Vec::with_capacity(n);
for _ in 0..n {
created.push(self.create(pool).await?);
}
Ok(created)
}
}
#[doc(hidden)]
pub mod __private {
use super::FactoryValue;
use std::marker::PhantomData;
pub struct Probe<T>(pub PhantomData<T>);
pub trait ViaValue<T> {
fn make(&self, field: &str, seq: u64) -> T;
fn make_key(&self, field: &str, seq: u64) -> T;
}
impl<T: FactoryValue> ViaValue<T> for &Probe<T> {
fn make(&self, field: &str, seq: u64) -> T {
T::factory_value(field, seq)
}
fn make_key(&self, field: &str, seq: u64) -> T {
T::factory_key(field, seq)
}
}
pub trait ViaDefault<T> {
fn make(&self, field: &str, seq: u64) -> T;
fn make_key(&self, field: &str, seq: u64) -> T;
}
impl<T: Default> ViaDefault<T> for Probe<T> {
fn make(&self, _field: &str, _seq: u64) -> T {
T::default()
}
fn make_key(&self, _field: &str, _seq: u64) -> T {
T::default()
}
}
}