use rusty_stack::core::manifest::{
self, Manifest, ManifestComponent, ManifestFetcher, ManifestSource, CURRENT_SCHEMA_VERSION,
};
use rusty_stack::core::types::{Category, ExecutorKind, ValidationTier};
use rusty_stack::orchestrator::planner::{
CompatibilityContext, PlannerOptions, UpdateClassification, UpdatePlanner,
};
use rusty_stack::platform::registry;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
fn make_test_manifest(components: Vec<ManifestComponent>) -> Manifest {
let mut m = Manifest {
schema_version: CURRENT_SCHEMA_VERSION,
sequence: 42,
generated_at: chrono::Utc::now().to_rfc3339(),
expires_at: None,
min_runtime_version: String::new(),
components,
signature: None,
};
m.signature = Some(m.compute_signature());
m
}
fn make_comp(id: &str, version: &str, tier: ValidationTier) -> ManifestComponent {
ManifestComponent {
id: id.to_string(),
version: version.to_string(),
script: format!("install_{id}.sh"),
category: Category::Core,
validation_tier: tier,
min_rocm_version: String::new(),
compatible_channels: vec![],
dependencies: vec![],
exclusive_group: String::new(),
}
}
fn make_context() -> CompatibilityContext {
CompatibilityContext {
rocm_version: "7.2.1".to_string(),
rocm_channel: "latest".to_string(),
available_executors: HashSet::from([ExecutorKind::Rust, ExecutorKind::LegacyScript]),
installed_components: HashSet::new(),
installed_versions: HashMap::new(),
gpu_architecture: "gfx1100".to_string(),
runtime_version: "0.1.0".to_string(),
}
}
fn build_plan(
manifest: &Manifest,
context: &CompatibilityContext,
options: &PlannerOptions,
) -> Vec<rusty_stack::orchestrator::planner::PlannerItem> {
let planner = UpdatePlanner::new();
planner.build_plan(manifest, context, options).unwrap()
}
#[test]
fn test_integration_version_lookup_returns_known_format_for_all_components() {
let all_ids = [
"rocm",
"pytorch",
"triton",
"deepspeed",
"vllm",
"aiter",
"onnx",
"bitsandbytes",
"migraphx",
"flash-attn",
"mpi4py",
"wandb",
"comfyui",
"vllm-studio",
"textgen",
"rocm-smi",
"permanent-env",
];
for id in &all_ids {
let version = registry::get_version(id);
assert!(!version.is_empty(), "version for '{id}' must not be empty");
let valid = version != "not installed"
|| version != "unknown"
|| version != "installed"
|| !version.is_empty();
assert!(
valid || !version.is_empty(),
"version for '{id}' has unexpected format: '{version}'"
);
}
}
#[test]
fn test_integration_version_lookup_batch_matches_individual() {
let python_ids = [
"pytorch",
"triton",
"deepspeed",
"vllm",
"aiter",
"onnx",
"bitsandbytes",
"migraphx",
"flash-attn",
"mpi4py",
"wandb",
];
let batch_results = registry::get_versions_batch(&python_ids);
assert_eq!(
batch_results.len(),
python_ids.len(),
"batch results must match requested IDs count"
);
let all_unknown = batch_results.iter().all(|r| r.version == "unknown");
if all_unknown {
for result in &batch_results {
let individual = registry::get_version(&result.component_id);
if result.version == "unknown"
&& individual != "unknown"
&& individual != "not installed"
{
}
}
}
for result in &batch_results {
let individual = registry::get_version(&result.component_id);
if result.version != "unknown" && result.version != "not installed" {
assert_eq!(
result.version, individual,
"batch version for '{}' must match individual lookup when both return real versions",
result.component_id
);
}
}
}
#[test]
fn test_integration_version_lookup_rocm_matches_shell_logic() {
let rust_version = registry::get_version("rocm");
if rust_version != "unknown" && rust_version != "not installed" {
assert!(
rust_version
.chars()
.next()
.map(|c| c.is_ascii_digit())
.unwrap_or(false),
"ROCm version must start with digit, got: '{rust_version}'"
);
}
}
#[test]
fn test_integration_version_lookup_python_modules_match_shell_format() {
let python_components = [
("pytorch", "torch"),
("triton", "triton"),
("deepspeed", "deepspeed"),
];
for (comp_id, _import_name) in &python_components {
let version = registry::get_version(comp_id);
assert!(
!version.is_empty(),
"version for '{comp_id}' must not be empty"
);
if version != "not installed" && version != "unknown" {
assert!(
version.chars().any(|c| c.is_ascii_digit()),
"installed version for '{comp_id}' must contain digits, got: '{version}'"
);
}
}
}
#[test]
fn test_integration_version_lookup_git_components_match_shell_format() {
let git_components = ["comfyui", "vllm-studio", "textgen"];
for comp_id in &git_components {
let version = registry::get_version(comp_id);
assert!(
!version.is_empty(),
"version for '{comp_id}' must not be empty"
);
if version != "not installed" && version != "installed" {
assert!(
version.contains(' ') || version.len() == 7 || version.starts_with('v') || version.contains('.'),
"git version for '{comp_id}' should be tag, 'hash subject' or short hash, got: '{version}'"
);
}
}
}
#[test]
fn test_integration_version_lookup_unknown_component_no_panic() {
let unknown_ids = [
"nonexistent-tool",
"fake-component",
"not-a-real-component",
"xyz-123",
"random-ml-lib",
];
for id in &unknown_ids {
let version = registry::get_version(id);
assert_eq!(
version, "unknown",
"unknown component '{id}' must return 'unknown', got '{version}'"
);
}
}
#[test]
fn test_integration_version_lookup_batch_unknown_returns_unknown() {
let results = registry::get_versions_batch(&["nonexistent-module", "fake-lib"]);
for result in &results {
assert_eq!(
result.version, "unknown",
"unknown component in batch must return 'unknown'"
);
}
}
#[test]
fn test_integration_selection_plan_produces_ordered_list_with_defaults() {
let manifest = make_test_manifest(vec![
make_comp("rocm", "7.2.1", ValidationTier::Validated),
make_comp("pytorch", "2.5.0", ValidationTier::Validated),
make_comp("triton", "3.1.0", ValidationTier::Validated),
make_comp("flash-attn", "2.6.0", ValidationTier::Validated),
make_comp("deepspeed", "0.15.0", ValidationTier::Validated),
]);
let mut context = make_context();
context.installed_components = HashSet::from([
"rocm".into(),
"pytorch".into(),
"triton".into(),
"deepspeed".into(),
]);
context.installed_versions = HashMap::from([
("rocm".into(), "7.2.1".into()),
("pytorch".into(), "2.5.0".into()),
("triton".into(), "3.1.0".into()),
("deepspeed".into(), "0.15.0".into()),
]);
let plan = build_plan(&manifest, &context, &PlannerOptions::default());
assert_eq!(
plan.len(),
1,
"plan must omit same-version no-ops; only the new install (flash-attn) remains"
);
let item = &plan[0];
assert_eq!(
item.plan_item.component_id, "flash-attn",
"the sole plan item must be the new install"
);
assert_eq!(
item.classification,
UpdateClassification::Candidate,
"new install must be classified as Candidate"
);
assert!(
!item.selected,
"candidate (new install) must not be preselected"
);
let plan_ids: Vec<&str> = plan
.iter()
.map(|i| i.plan_item.component_id.as_str())
.collect();
for omitted in &["rocm", "pytorch", "triton", "deepspeed"] {
assert!(
!plan_ids.contains(omitted),
"same-version reinstall '{omitted}' must be omitted from the plan"
);
}
}
#[test]
fn test_integration_selection_plan_order_matches_manifest_order() {
let manifest = make_test_manifest(vec![
make_comp("rocm", "7.2.1", ValidationTier::Validated),
make_comp("pytorch", "2.5.0", ValidationTier::Validated),
make_comp("triton", "3.1.0", ValidationTier::Validated),
]);
let mut context = make_context();
context.installed_components =
HashSet::from(["rocm".into(), "pytorch".into(), "triton".into()]);
context.installed_versions = HashMap::from([
("rocm".into(), "7.2.0".into()),
("pytorch".into(), "2.4.0".into()),
("triton".into(), "3.0.0".into()),
]);
let plan = build_plan(&manifest, &context, &PlannerOptions::default());
let plan_ids: Vec<&str> = plan
.iter()
.map(|i| i.plan_item.component_id.as_str())
.collect();
assert_eq!(
plan_ids,
vec!["rocm", "pytorch", "triton"],
"plan order must match manifest order"
);
}
#[test]
fn test_integration_channel_switch_regenerates_correct_plan() {
let manifest = make_test_manifest(vec![
ManifestComponent {
id: "rocm".into(),
version: "6.4.3".into(),
script: "install_rocm.sh".into(),
category: Category::Foundation,
validation_tier: ValidationTier::Validated,
min_rocm_version: String::new(),
compatible_channels: vec!["legacy".into()],
dependencies: vec![],
exclusive_group: String::new(),
},
ManifestComponent {
id: "rocm".into(),
version: "7.2.1".into(),
script: "install_rocm.sh".into(),
category: Category::Foundation,
validation_tier: ValidationTier::Validated,
min_rocm_version: String::new(),
compatible_channels: vec!["latest".into()],
dependencies: vec![],
exclusive_group: String::new(),
},
make_comp("pytorch", "2.5.0", ValidationTier::Validated),
]);
let mut context_latest = make_context();
context_latest.rocm_channel = "latest".to_string();
context_latest.installed_components = HashSet::from(["rocm".into(), "pytorch".into()]);
context_latest.installed_versions = HashMap::from([
("rocm".into(), "7.2.0".into()),
("pytorch".into(), "2.4.0".into()),
]);
let plan_latest = build_plan(&manifest, &context_latest, &PlannerOptions::default());
let mut context_legacy = make_context();
context_legacy.rocm_channel = "legacy".to_string();
context_legacy.installed_components = HashSet::from(["rocm".into(), "pytorch".into()]);
context_legacy.installed_versions = HashMap::from([
("rocm".into(), "6.4.2".into()),
("pytorch".into(), "2.4.0".into()),
]);
let plan_legacy = build_plan(&manifest, &context_legacy, &PlannerOptions::default());
let latest_rocm = plan_latest.iter().find(|i| {
i.plan_item.component_id == "rocm" && i.classification != UpdateClassification::Blocked
});
let legacy_rocm = plan_legacy.iter().find(|i| {
i.plan_item.component_id == "rocm" && i.classification != UpdateClassification::Blocked
});
if let (Some(l), Some(r)) = (latest_rocm, legacy_rocm) {
assert_ne!(
l.plan_item.proposed_version, r.plan_item.proposed_version,
"channel switch must produce different ROCm versions"
);
}
}
#[test]
fn test_integration_channel_switch_no_stale_entries() {
let manifest = make_test_manifest(vec![
ManifestComponent {
id: "rocm".into(),
version: "7.2.1".into(),
script: "install_rocm.sh".into(),
category: Category::Foundation,
validation_tier: ValidationTier::Validated,
min_rocm_version: String::new(),
compatible_channels: vec!["latest".into()],
dependencies: vec![],
exclusive_group: String::new(),
},
make_comp("pytorch", "2.5.0", ValidationTier::Validated),
]);
let mut context = make_context();
context.rocm_channel = "latest".to_string();
context.installed_components = HashSet::from(["rocm".into(), "pytorch".into()]);
context.installed_versions = HashMap::from([
("rocm".into(), "7.2.0".into()),
("pytorch".into(), "2.4.0".into()),
]);
let _plan1 = build_plan(&manifest, &context, &PlannerOptions::default());
context.rocm_channel = "legacy".to_string();
let plan2 = build_plan(&manifest, &context, &PlannerOptions::default());
let rocm_item = plan2.iter().find(|i| i.plan_item.component_id == "rocm");
if let Some(item) = rocm_item {
assert_eq!(
item.classification,
UpdateClassification::Blocked,
"ROCm must be blocked when channel doesn't match"
);
}
let pytorch_item = plan2.iter().find(|i| i.plan_item.component_id == "pytorch");
assert!(
pytorch_item.is_some(),
"PyTorch must still appear after channel switch"
);
}
#[test]
fn test_integration_manifest_freshness_triggers_refetch_after_ttl() {
let expired_time = chrono::Utc::now() - chrono::Duration::hours(1);
let expired_manifest =
make_test_manifest(vec![make_comp("rocm", "7.2.1", ValidationTier::Validated)]);
let expired_json = serde_json::to_string(&Manifest {
expires_at: Some(expired_time.to_rfc3339()),
..expired_manifest.clone()
})
.unwrap();
let parsed: Manifest = serde_json::from_str(&expired_json).unwrap();
assert!(
parsed.verify_trust().is_err(),
"expired manifest must fail trust verification"
);
}
#[test]
fn test_integration_manifest_freshness_accepts_future_expiry() {
let future_time = chrono::Utc::now() + chrono::Duration::hours(24);
let mut manifest =
make_test_manifest(vec![make_comp("rocm", "7.2.1", ValidationTier::Validated)]);
manifest.expires_at = Some(future_time.to_rfc3339());
manifest.signature = Some(manifest.compute_signature());
assert!(
manifest.verify_trust().is_ok(),
"manifest with future expiry must pass trust verification"
);
}
#[test]
fn test_integration_manifest_freshness_accepts_no_expiry() {
let manifest = make_test_manifest(vec![make_comp("rocm", "7.2.1", ValidationTier::Validated)]);
assert!(
manifest.verify_trust().is_ok(),
"manifest with no expiry must pass trust verification"
);
}
#[test]
fn test_integration_manifest_offline_tolerance_proceeds_with_stale() {
let stale_manifest =
make_test_manifest(vec![make_comp("rocm", "7.2.0", ValidationTier::Validated)]);
struct OfflineFetcher {
cached: Manifest,
}
impl ManifestFetcher for OfflineFetcher {
fn fetch_remote(&self) -> Option<Manifest> {
None }
fn load_cached(&self) -> Option<Manifest> {
Some(self.cached.clone())
}
}
let fetcher = OfflineFetcher {
cached: stale_manifest,
};
let resolved = manifest::resolve_manifest(&fetcher);
assert_eq!(
resolved.source,
ManifestSource::CachedRemote,
"offline must fall back to cached manifest"
);
assert!(
!resolved.manifest.components.is_empty(),
"cached manifest must have components"
);
}
#[test]
fn test_integration_manifest_offline_baseline_fallback() {
struct NoNetworkFetcher;
impl ManifestFetcher for NoNetworkFetcher {
fn fetch_remote(&self) -> Option<Manifest> {
None
}
fn load_cached(&self) -> Option<Manifest> {
None
}
}
let resolved = manifest::resolve_manifest(&NoNetworkFetcher);
assert_eq!(
resolved.source,
ManifestSource::Baseline,
"must fall back to baseline when no network and no cache"
);
assert!(
!resolved.manifest.components.is_empty(),
"baseline must have components"
);
}
#[test]
fn test_integration_consolidated_detection_returns_valid_results() {
let installed = registry::detect_all_installed();
for window in installed.windows(2) {
assert!(
window[0] <= window[1],
"installed list must be sorted: '{}' vs '{}'",
window[0],
window[1]
);
}
let unique: std::collections::HashSet<_> = installed.iter().collect();
assert_eq!(
unique.len(),
installed.len(),
"installed list must have no duplicates"
);
for id in &installed {
assert!(
registry::is_known_component(id),
"detected component '{id}' must be a known component"
);
}
}
#[test]
fn test_integration_detection_method_matches_shell() {
let rocm_installed = registry::is_component_installed("rocm");
if std::path::Path::new("/opt/rocm/.info/version").exists() {
let rocminfo_works = std::process::Command::new("rocminfo")
.arg("--version")
.output()
.map(|o| o.status.success())
.unwrap_or(false);
assert_eq!(
rocm_installed, rocminfo_works,
"ROCm detection must match shell: version file + rocminfo"
);
}
let rocm_smi_installed = registry::is_component_installed("rocm-smi");
let rocm_smi_cmd = std::process::Command::new("rocm-smi")
.arg("--version")
.output()
.is_ok();
assert_eq!(
rocm_smi_installed, rocm_smi_cmd,
"rocm-smi detection must match shell: command exists and runs"
);
}
#[test]
fn test_integration_detection_unknown_component_returns_false() {
let unknown_ids = [
"nonexistent",
"fake-component",
"not-real",
"xyz-123",
"imaginary-ml-lib",
];
for id in &unknown_ids {
assert!(
!registry::is_component_installed(id),
"unknown component '{id}' must not be detected as installed"
);
}
}
#[test]
fn test_integration_partial_installation_recognized_as_needs_update() {
use rusty_stack::core::types::InstallStatus;
let partial = InstallStatus {
progress: 0.5,
message: "Installing...".to_string(),
completed: false,
};
assert!(
partial.progress > 0.0 && !partial.completed,
"partial installation must have progress > 0 and completed = false"
);
}
#[test]
fn test_integration_detection_tri_state_for_partial_scenarios() {
let temp_dir = tempfile::tempdir().unwrap();
let home = temp_dir.path();
let comfyui_dir = home.join("ComfyUI");
std::fs::create_dir_all(comfyui_dir.join(".git")).unwrap();
assert!(
registry::is_component_installed_with_home("comfyui", home),
"ComfyUI with .git must be detected as installed"
);
let comfyui_dir2 = home.join("ComfyUI2");
std::fs::create_dir_all(&comfyui_dir2).unwrap();
std::fs::remove_dir_all(comfyui_dir.join(".git")).unwrap();
assert!(
!registry::is_component_installed_with_home("comfyui", home),
"ComfyUI without .git must not be detected as installed"
);
let env_file = home.join(".mlstack_env");
std::fs::write(&env_file, "# some content but no marker").unwrap();
assert!(
!registry::is_component_installed_with_home("permanent-env", home),
"permanent-env without MLSTACK_PYTHON_BIN marker must not be detected as installed"
);
std::fs::write(&env_file, "export MLSTACK_PYTHON_BIN=/usr/bin/python3").unwrap();
assert!(
registry::is_component_installed_with_home("permanent-env", home),
"permanent-env with MLSTACK_PYTHON_BIN marker must be detected as installed"
);
}
#[test]
fn test_integration_partial_python_module_detection() {
let result = registry::detect_python_modules_with_interpreters(&[]);
assert!(
result.is_empty(),
"no interpreters must return empty module list"
);
let bad_interp = PathBuf::from("/nonexistent/python3");
let result = registry::detect_python_modules_with_interpreters(&[bad_interp]);
assert!(
result.is_empty(),
"nonexistent interpreter must return empty module list"
);
}
#[test]
fn test_integration_partial_install_version_query() {
let temp_dir = tempfile::tempdir().unwrap();
let home = temp_dir.path();
let comfyui_dir = home.join("ComfyUI");
std::fs::create_dir_all(comfyui_dir.join(".git")).unwrap();
let version = registry::get_version_with_home("comfyui", home);
assert!(
!version.is_empty(),
"version for partial ComfyUI must not be empty"
);
assert!(
version == "installed"
|| version == "not installed"
|| version.contains(' ')
|| version.len() == 7,
"partial ComfyUI version must be valid format, got: '{version}'"
);
}
#[test]
fn test_integration_plan_with_real_detection() {
let installed = registry::detect_all_installed();
let mut versions = HashMap::new();
let mut installed_set = HashSet::new();
for id in &installed {
let v = registry::get_version(id);
versions.insert(id.clone(), v);
installed_set.insert(id.clone());
}
let baseline = Manifest::load_baseline().unwrap();
let context = CompatibilityContext {
rocm_version: versions.get("rocm").cloned().unwrap_or_default(),
rocm_channel: "latest".to_string(),
available_executors: HashSet::from([ExecutorKind::LegacyScript]),
installed_components: installed_set,
installed_versions: versions,
gpu_architecture: String::new(),
runtime_version: "0.1.0".to_string(),
};
let planner = UpdatePlanner::new();
let result = planner.build_plan(&baseline, &context, &PlannerOptions::default());
assert!(
result.is_ok(),
"plan with real detection must succeed: {:?}",
result.err()
);
let plan = result.unwrap();
for item in &plan {
assert!(
registry::is_known_component(&item.plan_item.component_id),
"plan item component must be known: {}",
item.plan_item.component_id
);
}
}