use crate::error::ShellError;
use crate::model::MarkovModel;
use std::path::Path;
pub fn sanitize_prefix(input: &str) -> Result<String, ShellError> {
let sanitized = input.replace('\0', "");
let trimmed = sanitized.trim();
if trimmed.is_empty() {
return Err(ShellError::InvalidInput {
message: "Empty prefix".into(),
});
}
if trimmed.len() < 2 {
return Err(ShellError::InvalidInput {
message: "Prefix too short (minimum 2 characters)".into(),
});
}
if trimmed.chars().any(|c| c.is_control() && c != '\t') {
return Err(ShellError::InvalidInput {
message: "Invalid control characters in input".into(),
});
}
Ok(trimmed.to_string())
}
pub fn load_model_graceful(path: &Path) -> Result<MarkovModel, ShellError> {
if !path.exists() {
return Err(ShellError::ModelNotFound {
path: path.to_path_buf(),
hint: "Run 'aprender-shell train' to create a model".into(),
});
}
match MarkovModel::load(path) {
Ok(model) => Ok(model),
Err(e) => {
let msg = e.to_string();
if msg.contains("Checksum")
|| msg.contains("checksum")
|| msg.contains("invalid")
|| msg.contains("corrupt")
{
Err(ShellError::ModelCorrupted {
path: path.to_path_buf(),
hint: "Model file is corrupted. Run 'aprender-shell train' to rebuild".into(),
})
} else {
Err(ShellError::ModelLoadFailed {
path: path.to_path_buf(),
cause: msg,
})
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[test]
fn test_empty_input_rejected() {
assert!(sanitize_prefix("").is_err());
let err = sanitize_prefix("").unwrap_err();
assert!(matches!(err, ShellError::InvalidInput { .. }));
}
#[test]
fn test_whitespace_only_rejected() {
assert!(sanitize_prefix(" ").is_err());
assert!(sanitize_prefix("\t\n").is_err());
assert!(sanitize_prefix(" \t ").is_err());
}
#[test]
fn test_null_bytes_removed() {
let result = sanitize_prefix("git \0status").expect("should succeed");
assert_eq!(result, "git status");
assert!(!result.contains('\0'));
}
#[test]
fn test_short_input_rejected() {
assert!(sanitize_prefix("g").is_err());
let err = sanitize_prefix("a").unwrap_err();
assert!(matches!(err, ShellError::InvalidInput { .. }));
}
#[test]
fn test_two_char_input_accepted() {
assert!(sanitize_prefix("gi").is_ok());
assert!(sanitize_prefix("ls").is_ok());
assert!(sanitize_prefix("cd").is_ok());
}
#[test]
fn test_double_dash_accepted() {
assert!(sanitize_prefix("--help").is_ok());
assert!(sanitize_prefix("-- filename").is_ok());
assert!(sanitize_prefix("git checkout --").is_ok());
}
#[test]
fn test_control_chars_rejected() {
assert!(sanitize_prefix("git\x07status").is_err());
assert!(sanitize_prefix("git\x1bstatus").is_err());
}
#[test]
fn test_tab_allowed() {
assert!(sanitize_prefix("git\tstatus").is_ok());
}
#[test]
fn test_valid_commands_accepted() {
assert!(sanitize_prefix("git status").is_ok());
assert!(sanitize_prefix("cargo build --release").is_ok());
assert!(sanitize_prefix("docker ps -a").is_ok());
assert!(sanitize_prefix("kubectl get pods").is_ok());
}
#[test]
fn test_whitespace_trimmed() {
let result = sanitize_prefix(" git status ").expect("should succeed");
assert_eq!(result, "git status");
}
#[test]
fn test_missing_model_graceful_error() {
let result = load_model_graceful(Path::new("/nonexistent/path/model.bin"));
assert!(matches!(result, Err(ShellError::ModelNotFound { .. })));
if let Err(ShellError::ModelNotFound { hint, .. }) = result {
assert!(hint.contains("train"));
}
}
#[test]
fn test_corrupted_model_graceful_error() {
let mut tmp = NamedTempFile::new().expect("create temp file");
tmp.write_all(b"GARBAGE DATA NOT A MODEL")
.expect("write garbage");
tmp.flush().expect("flush");
let result = load_model_graceful(tmp.path());
assert!(
matches!(
result,
Err(ShellError::ModelCorrupted { .. }) | Err(ShellError::ModelLoadFailed { .. })
),
"Expected graceful error for corrupted model"
);
}
#[test]
fn test_valid_model_loads_successfully() {
let mut model = MarkovModel::new(3);
model.train(&["git status".to_string(), "git commit".to_string()]);
let tmp = NamedTempFile::new().expect("create temp file");
model.save(tmp.path()).expect("save model");
let result = load_model_graceful(tmp.path());
assert!(result.is_ok(), "Valid model should load successfully");
}
#[test]
fn test_empty_file_graceful_error() {
let tmp = NamedTempFile::new().expect("create temp file");
let result = load_model_graceful(tmp.path());
assert!(result.is_err(), "Empty file should fail gracefully");
}
}