use crate::core::config::{ExtractionConfig, FileExtractionConfig};
use bytes::Bytes;
use std::path::PathBuf;
#[derive(Debug, Clone)]
pub enum ExtractionSource {
File { path: PathBuf, mime_hint: Option<String> },
Bytes { data: Bytes, mime_type: String },
}
#[derive(Debug, Clone)]
pub struct ExtractionRequest {
pub source: ExtractionSource,
pub config: ExtractionConfig,
pub file_overrides: Option<FileExtractionConfig>,
}
impl ExtractionRequest {
pub fn file(path: impl Into<PathBuf>, config: ExtractionConfig) -> Self {
Self {
source: ExtractionSource::File {
path: path.into(),
mime_hint: None,
},
config,
file_overrides: None,
}
}
pub fn file_with_mime(path: impl Into<PathBuf>, mime_hint: impl Into<String>, config: ExtractionConfig) -> Self {
Self {
source: ExtractionSource::File {
path: path.into(),
mime_hint: Some(mime_hint.into()),
},
config,
file_overrides: None,
}
}
pub fn bytes(data: impl Into<Bytes>, mime_type: impl Into<String>, config: ExtractionConfig) -> Self {
Self {
source: ExtractionSource::Bytes {
data: data.into(),
mime_type: mime_type.into(),
},
config,
file_overrides: None,
}
}
pub fn with_overrides(mut self, overrides: FileExtractionConfig) -> Self {
self.file_overrides = Some(overrides);
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn file_creates_file_source() {
let req = ExtractionRequest::file("/tmp/doc.pdf", ExtractionConfig::default());
match &req.source {
ExtractionSource::File { path, mime_hint } => {
assert_eq!(path, &PathBuf::from("/tmp/doc.pdf"));
assert!(mime_hint.is_none());
}
_ => panic!("expected File source"),
}
assert!(req.file_overrides.is_none());
}
#[test]
fn bytes_creates_bytes_source() {
let req = ExtractionRequest::bytes(b"hello".as_slice(), "text/plain", ExtractionConfig::default());
match &req.source {
ExtractionSource::Bytes { data, mime_type } => {
assert_eq!(data.as_ref(), b"hello");
assert_eq!(mime_type, "text/plain");
}
_ => panic!("expected Bytes source"),
}
}
#[test]
fn file_with_mime_sets_hint() {
let req = ExtractionRequest::file_with_mime("/tmp/doc.pdf", "application/pdf", ExtractionConfig::default());
match &req.source {
ExtractionSource::File { mime_hint, .. } => {
assert_eq!(mime_hint.as_deref(), Some("application/pdf"));
}
_ => panic!("expected File source"),
}
}
#[test]
fn with_overrides_sets_file_overrides() {
let overrides = FileExtractionConfig::default();
let req = ExtractionRequest::file("/tmp/doc.pdf", ExtractionConfig::default()).with_overrides(overrides);
assert!(req.file_overrides.is_some());
}
}