use std::sync::LazyLock; use k8s_openapi::api::coordination::v1::Lease; use kube::Client; use kube::api::Api; use rocket::async_test; use rocket::futures::FutureExt; use rocket::http::{ContentType, Header, Status}; use rocket::tokio; use rocket::tokio::sync::Mutex; static LOCK: LazyLock> = LazyLock::new(|| Mutex::new(())); async fn delete_lease(name: &str) { let client = Client::try_default().await.unwrap(); let leases: Api = Api::default_namespaced(client); let _ = kube::runtime::wait::delete::delete_and_finalize( leases, name, &Default::default(), ) .await; } async fn get_lease(name: &str) -> Result { let client = Client::try_default().await.unwrap(); let leases: Api = Api::default_namespaced(client); leases.get(name).await } #[async_test] async fn test_lock_v1_success() { super::setup(); let _lock = &*LOCK.lock().await; delete_lease("reboot-lock-default").await; let client = super::async_client().await; let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); } #[async_test] async fn test_lock_v1_custom_group() { super::setup(); delete_lease("reboot-lock-testgroup").await; let client = super::async_client().await; let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org&group=testgroup") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let lease = get_lease("reboot-lock-testgroup").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); } #[async_test] async fn test_lock_v1_conflict() { super::setup(); let _lock = &*LOCK.lock().await; delete_lease("reboot-lock-default").await; let client = super::async_client().await; let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test2.example.org&wait=false") .dispatch() .await; assert_eq!(response.status(), Status::Conflict); let want_msg = concat!( "Another system is already rebooting:", " Apply failed with 1 conflict:", " conflict with \"test1.example.org\":", " .spec.holderIdentity", "\n", ); assert_eq!(response.into_string().await.as_deref(), Some(want_msg)); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); } #[async_test] async fn test_lock_v1_conflict_wait() { super::setup(); let _lock = &*LOCK.lock().await; tracing::info!("Deleting existing lease"); delete_lease("reboot-lock-default").await; tracing::info!("Creating first lease"); let client = super::async_client().await; let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); let timer = std::time::Instant::now(); let _task = tokio::spawn(async { tokio::time::sleep(std::time::Duration::from_secs(1)) .then(|_| async { tracing::info!("Deleting first lease"); delete_lease("reboot-lock-default").await }) .await }); tracing::info!("Creating second lease"); let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test2.example.org") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let duration = timer.elapsed().as_millis(); assert!(duration > 1000 && duration < 2000); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test2.example.org") ); } #[test] fn test_lock_v1_no_header() { super::setup(); let client = super::client(); let response = client .post("/api/v1/lock") .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch(); assert_eq!(response.status(), Status::BadRequest); assert_eq!( response.into_string().as_deref(), Some("Invalid lock header\n") ); } #[test] fn test_lock_v1_no_data() { super::setup(); let client = super::client(); let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("") .dispatch(); assert_eq!(response.status(), Status::UnprocessableEntity); assert_eq!( response.into_string().as_deref(), Some("Error processing request:\nhostname: missing\n") ); } #[async_test] async fn test_unlock_v1_success() { super::setup(); let _lock = &*LOCK.lock().await; delete_lease("reboot-lock-default").await; let client = super::async_client().await; let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); let response = client .post("/api/v1/unlock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; let status = response.status(); assert_eq!(response.into_string().await, None); assert_eq!(status, Status::Ok); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!(lease.spec.unwrap().holder_identity, None); } #[async_test] async fn test_unlock_v1_not_locked() { super::setup(); let _lock = &*LOCK.lock().await; delete_lease("reboot-lock-default").await; let client = super::async_client().await; let response = client .post("/api/v1/unlock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; let status = response.status(); assert_eq!(response.into_string().await, None); assert_eq!(status, Status::Ok); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!(lease.spec.unwrap().holder_identity.as_deref(), None); } #[async_test] async fn test_unlock_v1_not_mine() { super::setup(); let _lock = &*LOCK.lock().await; delete_lease("reboot-lock-default").await; let client = super::async_client().await; let response = client .post("/api/v1/lock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch() .await; assert_eq!(response.status(), Status::Ok); assert_eq!(response.into_string().await, None); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); let response = client .post("/api/v1/unlock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("hostname=test2.example.org") .dispatch() .await; let status = response.status(); assert_eq!(response.into_string().await, None); assert_eq!(status, Status::Ok); let lease = get_lease("reboot-lock-default").await.unwrap(); assert_eq!( lease.spec.unwrap().holder_identity.as_deref(), Some("test1.example.org") ); } #[test] fn test_unlock_v1_no_header() { super::setup(); let client = super::client(); let response = client .post("/api/v1/unlock") .header(ContentType::Form) .body("hostname=test1.example.org") .dispatch(); assert_eq!(response.status(), Status::BadRequest); assert_eq!( response.into_string().as_deref(), Some("Invalid lock header\n") ); } #[test] fn test_unlock_v1_no_data() { super::setup(); let client = super::client(); let response = client .post("/api/v1/unlock") .header(Header::new("K8s-Reboot-Lock", "lock")) .header(ContentType::Form) .body("") .dispatch(); assert_eq!(response.status(), Status::UnprocessableEntity); assert_eq!( response.into_string().as_deref(), Some("Error processing request:\nhostname: missing\n") ); }