use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use serde_json::{Value as JsonValue, json};
use uuid::Uuid;
use chrono::{DateTime, Utc};
use crate::{TestResult, database::TestDatabase};
#[async_trait]
pub trait Factory<T: Send>: Send + Sync {
async fn create(&self) -> TestResult<T>;
async fn create_many(&self, count: usize) -> TestResult<Vec<T>> {
let mut results = Vec::with_capacity(count);
for _ in 0..count {
results.push(self.create().await?);
}
Ok(results)
}
fn build(&self) -> TestResult<T>;
fn build_many(&self, count: usize) -> TestResult<Vec<T>> {
let mut results = Vec::with_capacity(count);
for _ in 0..count {
results.push(self.build()?);
}
Ok(results)
}
}
#[derive(Clone)]
pub struct FactoryBuilder<T> {
attributes: HashMap<String, JsonValue>,
database: Option<Arc<TestDatabase>>,
_phantom: std::marker::PhantomData<T>,
}
impl<T> FactoryBuilder<T> {
pub fn new() -> Self {
Self {
attributes: HashMap::new(),
database: None,
_phantom: std::marker::PhantomData,
}
}
pub fn with<V: serde::Serialize>(mut self, key: &str, value: V) -> Self {
if let Ok(json_value) = serde_json::to_value(value) {
self.attributes.insert(key.to_string(), json_value);
}
self
}
pub fn with_attributes(mut self, attributes: HashMap<String, JsonValue>) -> Self {
self.attributes.extend(attributes);
self
}
pub fn with_database(mut self, database: Arc<TestDatabase>) -> Self {
self.database = Some(database);
self
}
pub fn with_relationship_data(mut self, name: &str, data: JsonValue) -> Self {
self.attributes.insert(format!("{}_data", name), data);
self
}
pub fn attributes(&self) -> &HashMap<String, JsonValue> {
&self.attributes
}
}
impl<T> Default for FactoryBuilder<T> {
fn default() -> Self {
Self::new()
}
}
pub trait HasId {
fn id(&self) -> JsonValue;
}
#[derive(Clone)]
pub struct UserFactory {
builder: FactoryBuilder<User>,
}
#[derive(Debug, Clone)]
pub struct User {
pub id: Uuid,
pub name: String,
pub email: String,
pub created_at: DateTime<Utc>,
pub updated_at: Option<DateTime<Utc>>,
}
impl HasId for User {
fn id(&self) -> JsonValue {
json!(self.id)
}
}
impl UserFactory {
pub fn new() -> Self {
let builder = FactoryBuilder::new();
Self { builder }
}
pub fn admin(self) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("role", "admin");
new_self
}
pub fn named(self, name: &str) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("name", name);
new_self
}
pub fn with_email(self, email: &str) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("email", email);
new_self
}
pub fn with_posts(self, count: usize) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("posts_count", count);
new_self
}
}
#[async_trait]
impl Factory<User> for UserFactory {
async fn create(&self) -> TestResult<User> {
let user = self.build()?;
if let Some(db) = &self.builder.database {
let insert_sql = r#"
INSERT INTO users (id, name, email, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5)
"#;
sqlx::query(insert_sql)
.bind(&user.id)
.bind(&user.name)
.bind(&user.email)
.bind(&user.created_at)
.bind(&user.updated_at)
.execute(db.pool())
.await?;
}
Ok(user)
}
fn build(&self) -> TestResult<User> {
let attrs = &self.builder.attributes;
let id = attrs.get("id")
.and_then(|v| v.as_str())
.and_then(|s| Uuid::parse_str(s).ok())
.unwrap_or_else(Uuid::new_v4);
let name = attrs.get("name")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| format!("Test User {}", crate::utils::random_string(None)));
let email = attrs.get("email")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| crate::utils::random_email());
Ok(User {
id,
name,
email,
created_at: attrs.get("created_at")
.and_then(|v| v.as_str())
.and_then(|s| DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc))
.unwrap_or_else(Utc::now),
updated_at: attrs.get("updated_at")
.and_then(|v| {
if v.is_null() {
None
} else {
v.as_str()
.and_then(|s| DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc))
}
}),
})
}
}
impl Default for UserFactory {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct PostFactory {
builder: FactoryBuilder<Post>,
}
#[derive(Debug, Clone)]
pub struct Post {
pub id: Uuid,
pub title: String,
pub content: String,
pub user_id: Uuid,
pub created_at: DateTime<Utc>,
pub updated_at: Option<DateTime<Utc>>,
}
impl HasId for Post {
fn id(&self) -> JsonValue {
json!(self.id)
}
}
impl PostFactory {
pub fn new() -> Self {
let builder = FactoryBuilder::new();
Self { builder }
}
pub fn with_title(self, title: &str) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("title", title);
new_self
}
pub fn with_content(self, content: &str) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("content", content);
new_self
}
pub fn for_user(self, user_id: Uuid) -> Self {
let mut new_self = self;
new_self.builder = new_self.builder.with("user_id", user_id);
new_self
}
pub fn with_user(self) -> Self {
let mut new_self = self;
let user_id = Uuid::new_v4();
new_self.builder = new_self.builder.with("user_id", user_id);
new_self
}
}
#[async_trait]
impl Factory<Post> for PostFactory {
async fn create(&self) -> TestResult<Post> {
let post = self.build()?;
if let Some(db) = &self.builder.database {
let insert_sql = r#"
INSERT INTO posts (id, title, content, user_id, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6)
"#;
sqlx::query(insert_sql)
.bind(&post.id)
.bind(&post.title)
.bind(&post.content)
.bind(&post.user_id)
.bind(&post.created_at)
.bind(&post.updated_at)
.execute(db.pool())
.await?;
}
Ok(post)
}
fn build(&self) -> TestResult<Post> {
let attrs = &self.builder.attributes;
let id = attrs.get("id")
.and_then(|v| v.as_str())
.and_then(|s| Uuid::parse_str(s).ok())
.unwrap_or_else(Uuid::new_v4);
let title = attrs.get("title")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| format!("Test Post {}", crate::utils::random_string(None)));
let user_id = attrs.get("user_id")
.and_then(|v| v.as_str())
.and_then(|s| Uuid::parse_str(s).ok())
.unwrap_or_else(Uuid::new_v4);
Ok(Post {
id,
title,
content: attrs.get("content")
.and_then(|v| v.as_str())
.unwrap_or("This is test content for the post.")
.to_string(),
user_id,
created_at: attrs.get("created_at")
.and_then(|v| v.as_str())
.and_then(|s| DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc))
.unwrap_or_else(Utc::now),
updated_at: attrs.get("updated_at")
.and_then(|v| {
if v.is_null() {
None
} else {
v.as_str()
.and_then(|s| DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&Utc))
}
}),
})
}
}
impl Default for PostFactory {
fn default() -> Self {
Self::new()
}
}
pub struct Sequence {
current: std::sync::atomic::AtomicUsize,
}
impl Sequence {
pub fn new() -> Self {
Self {
current: std::sync::atomic::AtomicUsize::new(0),
}
}
pub fn next(&self) -> usize {
self.current.fetch_add(1, std::sync::atomic::Ordering::SeqCst)
}
pub fn next_string(&self, prefix: &str) -> String {
format!("{}{}", prefix, self.next())
}
}
impl Default for Sequence {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_user_factory_build() -> TestResult<()> {
let factory = UserFactory::new();
let user = factory.build()?;
assert!(!user.name.is_empty());
assert!(user.email.contains("@"));
assert!(user.created_at <= Utc::now());
Ok(())
}
#[test]
fn test_user_factory_with_custom_attributes() -> TestResult<()> {
let factory = UserFactory::new()
.named("John Doe")
.with_email("john@example.com");
let user = factory.build()?;
assert_eq!(user.name, "John Doe");
assert_eq!(user.email, "john@example.com");
Ok(())
}
#[test]
fn test_post_factory_build() -> TestResult<()> {
let factory = PostFactory::new();
let post = factory.build()?;
assert!(!post.title.is_empty());
assert!(!post.content.is_empty());
assert!(post.created_at <= Utc::now());
Ok(())
}
#[test]
fn test_sequence() {
let seq = Sequence::new();
assert_eq!(seq.next(), 0);
assert_eq!(seq.next(), 1);
assert_eq!(seq.next_string("user"), "user2");
}
#[test]
fn test_factory_builder() {
let builder = FactoryBuilder::<User>::new()
.with("name", "Test User")
.with("email", "test@example.com");
assert_eq!(builder.attributes().get("name"), Some(&json!("Test User")));
assert_eq!(builder.attributes().get("email"), Some(&json!("test@example.com")));
}
#[tokio::test]
async fn test_factory_create_many() -> TestResult<()> {
let factory = UserFactory::new();
let users = factory.build_many(3)?;
assert_eq!(users.len(), 3);
for i in 0..users.len() {
for j in (i + 1)..users.len() {
assert_ne!(users[i].id, users[j].id);
}
}
Ok(())
}
}