[PATCH 3/3 v4] add KUnit coverage on Rust `ww_mutex` implementation

["Onur Özkan" <work@onurozkan.dev>]

Adds coverage around the core `ww_mutex` functionality

Signed-off-by: Onur Özkan <work@onurozkan.dev>
---
 rust/kernel/sync/lock/ww_mutex.rs | 116 ++++++++++++++++++++++++++++++
 1 file changed, 116 insertions(+)

diff --git a/rust/kernel/sync/lock/ww_mutex.rs b/rust/kernel/sync/lock/ww_mutex.rs
index 0f77e03304e4..30c6068f315d 100644
--- a/rust/kernel/sync/lock/ww_mutex.rs
+++ b/rust/kernel/sync/lock/ww_mutex.rs
@@ -11,6 +11,7 @@
 use crate::{bindings, str::CStr};
 use core::marker::PhantomData;
 use core::{cell::UnsafeCell, pin::Pin};
+use macros::kunit_tests;
 use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
 
 /// Create static `WwClass` instances.
@@ -417,3 +418,118 @@ fn drop(&mut self) {
         unsafe { bindings::ww_mutex_unlock(self.mutex.as_ptr()) };
     }
 }
+
+#[kunit_tests(rust_kernel_ww_mutex)]
+mod tests {
+    use crate::alloc::KBox;
+    use crate::c_str;
+    use crate::prelude::*;
+    use pin_init::stack_pin_init;
+
+    use super::*;
+
+    // A simple coverage on `define_ww_class` macro.
+    define_ww_class!(TEST_WOUND_WAIT_CLASS, wound_wait, c_str!("test_wound_wait"));
+    define_ww_class!(TEST_WAIT_DIE_CLASS, wait_die, c_str!("test_wait_die"));
+
+    #[test]
+    fn test_ww_mutex_basic_lock_unlock() {
+        stack_pin_init!(let class = WwClass::new_wound_wait(c_str!("test_mutex_class")));
+
+        stack_pin_init!(let mutex = WwMutex::new(42, &class));
+
+        // Lock without context
+        let guard = mutex.lock(None).unwrap();
+        assert_eq!(*guard, 42);
+
+        // Drop the lock
+        drop(guard);
+
+        // Lock it again
+        let mut guard = mutex.lock(None).unwrap();
+        *guard = 100;
+        assert_eq!(*guard, 100);
+    }
+
+    #[test]
+    fn test_ww_mutex_trylock() {
+        stack_pin_init!(let class = WwClass::new_wound_wait(c_str!("trylock_class")));
+
+        stack_pin_init!(let mutex = WwMutex::new(123, &class));
+
+        // trylock on unlocked mutex should succeed
+        let guard = mutex.try_lock(None).unwrap();
+        assert_eq!(*guard, 123);
+        drop(guard);
+
+        // lock it first
+        let _guard1 = mutex.lock(None).unwrap();
+
+        // trylock should fail with EBUSY when already locked
+        let result = mutex.try_lock(None);
+        match result {
+            Err(e) => assert_eq!(e, EBUSY),
+            Ok(_) => panic!("Expected `EBUSY` but got success"),
+        }
+    }
+
+    #[test]
+    fn test_ww_mutex_is_locked() {
+        stack_pin_init!(let class = WwClass::new_wait_die(c_str!("locked_check_class")));
+
+        stack_pin_init!(let mutex = WwMutex::new("hello", &class));
+
+        // should not be locked initially
+        assert!(!mutex.is_locked());
+
+        let guard = mutex.lock(None).unwrap();
+        assert!(mutex.is_locked());
+
+        drop(guard);
+        assert!(!mutex.is_locked());
+    }
+
+    #[test]
+    fn test_ww_acquire_context() {
+        stack_pin_init!(let class = WwClass::new_wound_wait(c_str!("ctx_class")));
+
+        stack_pin_init!(let mutex1 = WwMutex::new(1, &class));
+        stack_pin_init!(let mutex2 = WwMutex::new(2, &class));
+
+        let mut ctx = KBox::pin_init(WwAcquireCtx::new(&class), GFP_KERNEL).unwrap();
+
+        // acquire multiple mutexes with same context
+        let guard1 = mutex1.lock(Some(&ctx)).unwrap();
+        let guard2 = mutex2.lock(Some(&ctx)).unwrap();
+
+        assert_eq!(*guard1, 1);
+        assert_eq!(*guard2, 2);
+
+        ctx.as_mut().done();
+
+        // we shouldn't be able to lock once it's `done`.
+        assert!(mutex1.lock(Some(&ctx)).is_err());
+        assert!(mutex2.lock(Some(&ctx)).is_err());
+    }
+
+    #[test]
+    fn test_with_global_classes() {
+        stack_pin_init!(let wound_wait_mutex = WwMutex::new(100, &TEST_WOUND_WAIT_CLASS));
+        stack_pin_init!(let wait_die_mutex = WwMutex::new(200, &TEST_WAIT_DIE_CLASS));
+
+        let ww_guard = wound_wait_mutex.lock(None).unwrap();
+        let wd_guard = wait_die_mutex.lock(None).unwrap();
+
+        assert_eq!(*ww_guard, 100);
+        assert_eq!(*wd_guard, 200);
+
+        assert!(wound_wait_mutex.is_locked());
+        assert!(wait_die_mutex.is_locked());
+
+        drop(ww_guard);
+        drop(wd_guard);
+
+        assert!(!wound_wait_mutex.is_locked());
+        assert!(!wait_die_mutex.is_locked());
+    }
+}
-- 
2.49.0