use dbx_core::agent_manager::{ AgentManager, AgentRegistry, ArtifactInfo, DriverInfo, InstalledDriver, JavaRuntimeConfig, JavaRuntimeMode, JreInfo, DEFAULT_JRE_KEY, }; use dbx_core::agent_service::{ build_agent_list, clear_agent_download_cache, github_url_to_r2_path, import_agent_jar, import_agents_from_zip, is_app_version_compatible, jre_needs_install, local_agent_jar_candidates, replace_download, uninstall_agent_driver, AgentProgressEvent, }; fn test_manager(name: &str) -> AgentManager { let dir = std::env::temp_dir().join(format!("dbx-agent-service-{name}-{}", uuid::Uuid::new_v4())); AgentManager::new_with_base_dir(dir) } fn registry_with_driver(db_type: &str, version: &str, jre: &str) -> AgentRegistry { let mut drivers = std::collections::HashMap::new(); drivers.insert( db_type.to_string(), DriverInfo { version: version.to_string(), label: db_type.to_string(), min_app_version: "0.1.0".to_string(), jre: jre.to_string(), jar: Some(ArtifactInfo { url: format!("https://example.com/dbx-agent-{db_type}.jar"), size: 42 }), native: std::collections::HashMap::new(), }, ); AgentRegistry { jre: None, jres: std::collections::HashMap::new(), drivers } } fn registry_with_jre_driver(db_type: &str, driver_version: &str, jre: &str, jre_version: &str) -> AgentRegistry { let mut registry = registry_with_driver(db_type, driver_version, jre); registry.jres.insert( jre.to_string(), JreInfo { version: jre_version.to_string(), platforms: std::collections::HashMap::new() }, ); registry } fn registry_with_native_driver(db_type: &str, version: &str, jre: &str) -> AgentRegistry { let mut drivers = std::collections::HashMap::new(); drivers.insert( db_type.to_string(), DriverInfo { version: version.to_string(), label: db_type.to_string(), min_app_version: "0.1.0".to_string(), jre: jre.to_string(), jar: Some(ArtifactInfo { url: format!("https://example.com/dbx-agent-{db_type}-legacy-placeholder.jar"), size: 0, }), native: [( AgentManager::current_platform().to_string(), ArtifactInfo { url: format!("https://example.com/dbx-agent-{db_type}"), size: 42 }, )] .into_iter() .collect(), }, ); AgentRegistry { jre: None, jres: std::collections::HashMap::new(), drivers } } #[test] fn built_in_agent_list_includes_expected_driver_labels() { let manager = test_manager("labels"); let agents = build_agent_list(&manager, None); assert!(agents.iter().any(|agent| agent.db_type == "tdengine" && agent.label == "TDengine")); assert!(agents.iter().any(|agent| agent.db_type == "iotdb" && agent.label == "Apache IoTDB")); assert!(agents.iter().any(|agent| agent.db_type == "yashandb" && agent.label == "崖山 YashanDB")); assert!(agents.iter().any(|agent| agent.db_type == "access" && agent.label == "Microsoft Access")); } #[test] fn agent_list_marks_installed_driver_update_when_registry_version_differs() { let manager = test_manager("update"); let jar_path = manager.driver_jar_path("h2"); std::fs::create_dir_all(jar_path.parent().unwrap()).unwrap(); write_test_agent_jar(&jar_path); manager .save_state(&dbx_core::agent_manager::AgentState { installed_drivers: [( "h2".to_string(), InstalledDriver { version: "0.1.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); let registry = registry_with_driver("h2", "0.2.0", "21"); let agents = build_agent_list(&manager, Some(®istry)); let h2 = agents.iter().find(|agent| agent.db_type == "h2").unwrap(); assert!(h2.installed); assert_eq!(h2.installed_version.as_deref(), Some("0.1.0")); assert_eq!(h2.version, "0.2.0"); assert_eq!(h2.size, 42); assert_eq!(h2.jre, "21"); assert!(h2.requires_java_runtime); assert!(h2.update_available); } #[test] fn agent_list_marks_update_when_installed_managed_jre_version_differs() { let manager = test_manager("jre-update"); let jar_path = manager.driver_jar_path("h2"); let java_path = manager.jre_java_path(DEFAULT_JRE_KEY); std::fs::create_dir_all(jar_path.parent().unwrap()).unwrap(); std::fs::create_dir_all(java_path.parent().unwrap()).unwrap(); write_test_agent_jar(&jar_path); std::fs::write(&java_path, b"java").unwrap(); manager .save_state(&dbx_core::agent_manager::AgentState { jre_versions: [(DEFAULT_JRE_KEY.to_string(), "21.0.10".to_string())].into_iter().collect(), installed_drivers: [( "h2".to_string(), InstalledDriver { version: "0.2.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); let registry = registry_with_jre_driver("h2", "0.2.0", DEFAULT_JRE_KEY, "21.0.11"); let agents = build_agent_list(&manager, Some(®istry)); let h2 = agents.iter().find(|agent| agent.db_type == "h2").unwrap(); assert!(h2.update_available); } #[test] fn agent_list_does_not_mark_jre_update_for_system_java_runtime() { let manager = test_manager("system-java-no-jre-update"); let jar_path = manager.driver_jar_path("dameng"); std::fs::create_dir_all(jar_path.parent().unwrap()).unwrap(); write_test_agent_jar(&jar_path); manager .save_state(&dbx_core::agent_manager::AgentState { java_runtime: JavaRuntimeConfig { mode: JavaRuntimeMode::System, custom_java_path: None }, installed_drivers: [( "dameng".to_string(), InstalledDriver { version: "0.2.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); let registry = registry_with_jre_driver("dameng", "0.2.0", DEFAULT_JRE_KEY, "21.0.11"); let agents = build_agent_list(&manager, Some(®istry)); let dameng = agents.iter().find(|agent| agent.db_type == "dameng").unwrap(); assert!(dameng.installed); assert!(!dameng.jre_installed); assert!(!dameng.update_available); } #[test] fn agent_list_does_not_require_jre_for_native_agent() { let manager = test_manager("native-no-jre"); let native_path = manager.driver_native_path("dameng"); std::fs::create_dir_all(native_path.parent().unwrap()).unwrap(); std::fs::write(&native_path, b"agent").unwrap(); manager .save_state(&dbx_core::agent_manager::AgentState { installed_drivers: [( "dameng".to_string(), InstalledDriver { version: "0.2.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); let agents = build_agent_list(&manager, None); let dameng = agents.iter().find(|agent| agent.db_type == "dameng").unwrap(); assert!(dameng.installed); assert!(!dameng.requires_java_runtime); assert!(dameng.jre_installed); } #[test] fn agent_list_does_not_require_jre_for_registry_native_agent() { let manager = test_manager("registry-native-no-jre"); let registry = registry_with_native_driver("xugu", "0.2.0", DEFAULT_JRE_KEY); let agents = build_agent_list(&manager, Some(®istry)); let xugu = agents.iter().find(|agent| agent.db_type == "xugu").unwrap(); assert!(!xugu.installed); assert_eq!(xugu.jre, DEFAULT_JRE_KEY); assert!(!xugu.requires_java_runtime); assert!(xugu.jre_installed); } #[test] fn agent_list_keeps_jre_requirement_for_installed_jar_when_registry_has_native() { let manager = test_manager("installed-jar-registry-native"); let jar_path = manager.driver_jar_path("xugu"); std::fs::create_dir_all(jar_path.parent().unwrap()).unwrap(); write_test_agent_jar(&jar_path); manager .save_state(&dbx_core::agent_manager::AgentState { installed_drivers: [( "xugu".to_string(), InstalledDriver { version: "0.2.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); let registry = registry_with_native_driver("xugu", "0.2.0", DEFAULT_JRE_KEY); let agents = build_agent_list(&manager, Some(®istry)); let xugu = agents.iter().find(|agent| agent.db_type == "xugu").unwrap(); assert!(xugu.installed); assert!(xugu.requires_java_runtime); assert!(!xugu.jre_installed); } #[test] fn agent_list_reports_missing_jre_for_uninstalled_java_drivers() { let manager = test_manager("uninstalled-java-driver-missing-jre"); let registry = registry_with_jre_driver("dameng", "0.2.0", DEFAULT_JRE_KEY, "21.0.11"); let agents = build_agent_list(&manager, Some(®istry)); let dameng = agents.iter().find(|agent| agent.db_type == "dameng").unwrap(); assert!(!dameng.installed); assert!(!dameng.jre_installed); } #[test] fn agent_list_uses_legacy_default_jre_version_when_checking_updates() { let manager = test_manager("legacy-jre-version"); let jar_path = manager.driver_jar_path("dameng"); let java_path = manager.jre_java_path(DEFAULT_JRE_KEY); std::fs::create_dir_all(jar_path.parent().unwrap()).unwrap(); std::fs::create_dir_all(java_path.parent().unwrap()).unwrap(); write_test_agent_jar(&jar_path); std::fs::write(&java_path, b"java").unwrap(); manager .save_state(&dbx_core::agent_manager::AgentState { jre_version: Some("21.0.11".to_string()), installed_drivers: [( "dameng".to_string(), InstalledDriver { version: "0.2.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); let registry = registry_with_jre_driver("dameng", "0.2.0", DEFAULT_JRE_KEY, "21.0.11"); let agents = build_agent_list(&manager, Some(®istry)); let dameng = agents.iter().find(|agent| agent.db_type == "dameng").unwrap(); assert!(!dameng.update_available); } #[test] fn jre_needs_install_when_managed_runtime_version_differs() { let manager = test_manager("jre-needs-install"); let java_path = manager.jre_java_path(DEFAULT_JRE_KEY); std::fs::create_dir_all(java_path.parent().unwrap()).unwrap(); std::fs::write(&java_path, b"java").unwrap(); manager .save_state(&dbx_core::agent_manager::AgentState { jre_versions: [(DEFAULT_JRE_KEY.to_string(), "21.0.10".to_string())].into_iter().collect(), ..Default::default() }) .unwrap(); let registry = registry_with_jre_driver("h2", "0.2.0", DEFAULT_JRE_KEY, "21.0.11"); assert!(jre_needs_install(&manager, ®istry, DEFAULT_JRE_KEY)); } #[test] fn local_agent_jar_candidates_include_monorepo_and_legacy_build_output() { let candidates = local_agent_jar_candidates("tdengine"); assert!(candidates.iter().any(|path| path.ends_with("agents/drivers/tdengine/build/libs/dbx-agent-tdengine.jar"))); assert!(candidates.iter().any(|path| path.ends_with("dbx-agents/tdengine/build/libs/dbx-agent-tdengine.jar"))); } #[test] fn github_agent_asset_urls_map_to_r2_paths_by_category() { assert_eq!( github_url_to_r2_path("https://github.com/t8y2/dbx-agents/releases/download/v1/dbx-jre-21.tar.gz", "jre"), "agents/jre/dbx-jre-21.tar.gz" ); assert_eq!( github_url_to_r2_path("https://github.com/t8y2/dbx-agents/releases/download/v1/dbx-agent-h2.jar", "driver"), "agents/drivers/dbx-agent-h2.jar" ); assert_eq!( github_url_to_r2_path("https://github.com/t8y2/dbx/releases/download/agents-v0.3.0/dbx-agent-h2.jar", "driver"), "agents/drivers/dbx-agent-h2.jar" ); } fn test_path(name: &str) -> std::path::PathBuf { std::env::temp_dir().join(format!("dbx-agent-service-{name}-{}", uuid::Uuid::new_v4())) } #[test] fn accepts_current_app_when_min_version_is_not_newer() { assert!(is_app_version_compatible("0.5.13", "0.5.13")); assert!(is_app_version_compatible("0.5.12", "0.5.13")); assert!(!is_app_version_compatible("0.5.14", "0.5.13")); } #[test] fn atomic_replace_moves_download_into_place() { let dir = test_path("atomic"); std::fs::create_dir_all(&dir).unwrap(); let dest = dir.join("agent.jar"); let tmp = dir.join("agent.jar.download"); std::fs::write(&dest, b"old").unwrap(); std::fs::write(&tmp, b"new").unwrap(); replace_download(&tmp, &dest).unwrap(); assert_eq!(std::fs::read(&dest).unwrap(), b"new"); assert!(!tmp.exists()); std::fs::remove_dir_all(dir).ok(); } #[test] fn agent_progress_event_serializes_backward_compatible_fields() { let event = AgentProgressEvent::transfer("driver", 512, 1024).with_batch(Some("h2"), Some(1), Some(2)); let value = serde_json::to_value(event).unwrap(); assert_eq!(value["step"], "driver"); assert_eq!(value["downloaded"], 512); assert_eq!(value["total"], 1024); assert_eq!(value["db_type"], "h2"); assert_eq!(value["current"], 1); assert_eq!(value["total_drivers"], 2); } #[test] fn local_jar_import_updates_driver_state() { let manager = test_manager("local-import"); let source = test_path("local-import-source").join("dbx-agent-h2.jar"); std::fs::create_dir_all(source.parent().unwrap()).unwrap(); write_test_agent_jar(&source); import_agent_jar(&manager, "h2", &source).unwrap(); assert_eq!(std::fs::read(manager.driver_jar_path("h2")).unwrap(), std::fs::read(&source).unwrap()); let state = manager.load_state(); let installed = state.installed_drivers.get("h2").unwrap(); assert_eq!(installed.version, "0.1.0-local"); assert_eq!(installed.jre, DEFAULT_JRE_KEY); } #[test] fn local_jar_import_rejects_corrupt_jar() { let manager = test_manager("local-import-corrupt"); let source = test_path("local-import-corrupt-source").join("dbx-agent-h2.jar"); std::fs::create_dir_all(source.parent().unwrap()).unwrap(); std::fs::write(&source, b"jar").unwrap(); let err = import_agent_jar(&manager, "h2", &source).unwrap_err(); assert!(err.contains("invalid or corrupt")); assert!(!manager.driver_jar_path("h2").exists()); assert!(!manager.load_state().installed_drivers.contains_key("h2")); } #[tokio::test] async fn uninstall_driver_removes_artifact_and_state() { let manager = test_manager("uninstall"); let jar_path = manager.driver_jar_path("h2"); std::fs::create_dir_all(jar_path.parent().unwrap()).unwrap(); std::fs::write(&jar_path, b"jar").unwrap(); let cache_dir = manager.download_cache_dir(); std::fs::create_dir_all(&cache_dir).unwrap(); let h2_cache = cache_dir.join("driver-h2-0.1.0-abc-agent.jar"); let dameng_cache = cache_dir.join("driver-dameng-0.1.0-abc-agent.jar"); let jre_cache = cache_dir.join("jre-21-21.0.11-abc-jre-download.tar.gz"); std::fs::write(&h2_cache, b"h2").unwrap(); std::fs::write(&dameng_cache, b"dameng").unwrap(); std::fs::write(&jre_cache, b"jre").unwrap(); manager .save_state(&dbx_core::agent_manager::AgentState { installed_drivers: [( "h2".to_string(), InstalledDriver { version: "0.1.0".to_string(), installed_at: "2026-05-18T00:00:00Z".to_string(), jre: DEFAULT_JRE_KEY.to_string(), }, )] .into_iter() .collect(), ..Default::default() }) .unwrap(); uninstall_agent_driver(&manager, "h2").await.unwrap(); assert!(!jar_path.exists()); assert!(!h2_cache.exists()); assert!(dameng_cache.exists()); assert!(jre_cache.exists()); assert!(!manager.load_state().installed_drivers.contains_key("h2")); } #[test] fn clear_download_cache_removes_only_cache_entries() { let manager = test_manager("clear-download-cache"); let cache_dir = manager.download_cache_dir(); let driver_dir = manager.driver_dir("h2"); std::fs::create_dir_all(&cache_dir).unwrap(); std::fs::create_dir_all(&driver_dir).unwrap(); let driver_cache = cache_dir.join("driver-h2-0.1.0-abc-agent.jar"); let jre_cache = cache_dir.join("jre-21-21.0.11-abc-jre-download.tar.gz"); let nested_cache = cache_dir.join("stale-dir"); let installed_driver = driver_dir.join("agent.jar"); std::fs::write(&driver_cache, b"h2").unwrap(); std::fs::write(&jre_cache, b"jre").unwrap(); std::fs::create_dir_all(&nested_cache).unwrap(); std::fs::write(nested_cache.join("artifact"), b"x").unwrap(); std::fs::write(&installed_driver, b"driver").unwrap(); clear_agent_download_cache(&manager).unwrap(); assert!(!driver_cache.exists()); assert!(!jre_cache.exists()); assert!(!nested_cache.exists()); assert!(installed_driver.exists()); } #[test] fn offline_zip_import_emits_progress_and_updates_state() { let manager = test_manager("offline-progress"); let zip_path = test_path("offline-progress-zip").join("agents.zip"); std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap(); write_offline_driver_zip(&zip_path, "h2", "0.2.0"); let events = std::sync::Mutex::new(Vec::new()); let result = import_agents_from_zip(&manager, &zip_path, |event| { events.lock().unwrap().push(event); }) .unwrap(); assert_eq!(result.drivers_installed, vec!["h2"]); let installed_jar = std::fs::read(manager.driver_jar_path("h2")).unwrap(); assert!(installed_jar.windows(b"Main-Class:".len()).any(|window| window == b"Main-Class:")); assert_eq!(manager.load_state().installed_drivers.get("h2").unwrap().version, "0.2.0"); let events = events.lock().unwrap(); assert!(events.iter().any(|event| event.step == "driver" && event.db_type.as_deref() == Some("H2"))); } #[test] fn offline_zip_import_rejects_corrupt_jar() { let manager = test_manager("offline-corrupt-driver"); let zip_path = test_path("offline-corrupt-driver-zip").join("agents.zip"); std::fs::create_dir_all(zip_path.parent().unwrap()).unwrap(); write_offline_driver_zip_with_jar(&zip_path, "h2", "0.2.0", b"jar".to_vec()); let err = import_agents_from_zip(&manager, &zip_path, |_| {}).unwrap_err(); assert!(err.contains("invalid or corrupt")); assert!(!manager.driver_jar_path("h2").exists()); assert!(!manager.load_state().installed_drivers.contains_key("h2")); } fn write_offline_driver_zip(path: &std::path::Path, db_type: &str, version: &str) { write_offline_driver_zip_with_jar(path, db_type, version, test_agent_jar_bytes()); } fn write_offline_driver_zip_with_jar(path: &std::path::Path, db_type: &str, version: &str, jar: Vec) { let file = std::fs::File::create(path).unwrap(); let mut zip = zip::ZipWriter::new(file); let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); let registry = serde_json::json!({ "drivers": { db_type: { "version": version, "label": db_type, "min_app_version": "0.1.0", "jre": DEFAULT_JRE_KEY, "jar": { "url": format!("https://example.com/dbx-agent-{db_type}.jar"), "size": jar.len() } } } }); zip.start_file("agent-registry.json", options).unwrap(); std::io::Write::write_all(&mut zip, registry.to_string().as_bytes()).unwrap(); zip.start_file(format!("drivers/dbx-agent-{db_type}.jar"), options).unwrap(); std::io::Write::write_all(&mut zip, &jar).unwrap(); zip.finish().unwrap(); } fn write_test_agent_jar(path: &std::path::Path) { std::fs::write(path, test_agent_jar_bytes()).unwrap(); } fn test_agent_jar_bytes() -> Vec { let cursor = std::io::Cursor::new(Vec::new()); let mut zip = zip::ZipWriter::new(cursor); let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored); // The production validator rejects corrupt driver artifacts by requiring a real JAR manifest. zip.start_file("META-INF/MANIFEST.MF", options).unwrap(); std::io::Write::write_all(&mut zip, b"Manifest-Version: 1.0\nMain-Class: com.dbx.agent.TestAgent\n\n").unwrap(); zip.finish().unwrap().into_inner() }