rustytorch/demos/rtx-model-zoo/src/error.rs

use thiserror::Error;

#[derive(Error, Debug, Clone, PartialEq)]
pub enum ModelZooError {
    #[error("Model not found: {model_id}")]
    ModelNotFound { model_id: String },

    #[error("Model already loaded: {model_id}")]
    ModelAlreadyLoaded { model_id: String },

    #[error("Model not loaded: {model_id}")]
    ModelNotLoaded { model_id: String },

    #[error("Model download failed: {model_id}, reason: {reason}")]
    DownloadFailed { model_id: String, reason: String },

    #[error("Invalid input data: {reason}")]
    InvalidInput { reason: String },

    #[error("Inference failed: {reason}")]
    InferenceFailed { reason: String },

    #[error("Configuration error: {message}")]
    ConfigError { message: String },

    #[error("IO error: {message}")]
    IoError { message: String },
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model_not_found_error() {
        let error = ModelZooError::ModelNotFound {
            model_id: "resnet50".to_string(),
        };
        assert_eq!(error.to_string(), "Model not found: resnet50");
    }

    #[test]
    fn test_model_already_loaded_error() {
        let error = ModelZooError::ModelAlreadyLoaded {
            model_id: "resnet50".to_string(),
        };
        assert_eq!(error.to_string(), "Model already loaded: resnet50");
    }

    #[test]
    fn test_model_not_loaded_error() {
        let error = ModelZooError::ModelNotLoaded {
            model_id: "resnet50".to_string(),
        };
        assert_eq!(error.to_string(), "Model not loaded: resnet50");
    }

    #[test]
    fn test_download_failed_error() {
        let error = ModelZooError::DownloadFailed {
            model_id: "resnet50".to_string(),
            reason: "Network timeout".to_string(),
        };
        assert_eq!(
            error.to_string(),
            "Model download failed: resnet50, reason: Network timeout"
        );
    }

    #[test]
    fn test_invalid_input_error() {
        let error = ModelZooError::InvalidInput {
            reason: "Invalid base64 encoding".to_string(),
        };
        assert_eq!(
            error.to_string(),
            "Invalid input data: Invalid base64 encoding"
        );
    }

    #[test]
    fn test_inference_failed_error() {
        let error = ModelZooError::InferenceFailed {
            reason: "Out of memory".to_string(),
        };
        assert_eq!(error.to_string(), "Inference failed: Out of memory");
    }

    #[test]
    fn test_config_error() {
        let error = ModelZooError::ConfigError {
            message: "Invalid download path".to_string(),
        };
        assert_eq!(
            error.to_string(),
            "Configuration error: Invalid download path"
        );
    }

    #[test]
    fn test_io_error() {
        let error = ModelZooError::IoError {
            message: "Failed to create directory".to_string(),
        };
        assert_eq!(error.to_string(), "IO error: Failed to create directory");
    }

    #[test]
    fn test_error_equality() {
        let error1 = ModelZooError::ModelNotFound {
            model_id: "resnet50".to_string(),
        };
        let error2 = ModelZooError::ModelNotFound {
            model_id: "resnet50".to_string(),
        };
        let error3 = ModelZooError::ModelNotFound {
            model_id: "yolov8".to_string(),
        };
        assert_eq!(error1, error2);
        assert_ne!(error1, error3);
    }

    #[test]
    fn test_error_clone() {
        let error = ModelZooError::ModelNotFound {
            model_id: "resnet50".to_string(),
        };
        let cloned = error.clone();
        assert_eq!(error, cloned);
    }

    #[test]
    fn test_error_debug() {
        let error = ModelZooError::ModelNotFound {
            model_id: "resnet50".to_string(),
        };
        let debug_str = format!("{error:?}");
        assert!(debug_str.contains("ModelNotFound"));
        assert!(debug_str.contains("resnet50"));
    }

    #[test]
    fn test_all_error_variants_different() {
        let errors = vec![
            ModelZooError::ModelNotFound {
                model_id: "test".to_string(),
            },
            ModelZooError::ModelAlreadyLoaded {
                model_id: "test".to_string(),
            },
            ModelZooError::ModelNotLoaded {
                model_id: "test".to_string(),
            },
            ModelZooError::DownloadFailed {
                model_id: "test".to_string(),
                reason: "reason".to_string(),
            },
            ModelZooError::InvalidInput {
                reason: "reason".to_string(),
            },
            ModelZooError::InferenceFailed {
                reason: "reason".to_string(),
            },
            ModelZooError::ConfigError {
                message: "message".to_string(),
            },
            ModelZooError::IoError {
                message: "message".to_string(),
            },
        ];
        assert_eq!(errors.len(), 8);
    }
}