f - Catch async KVStore test panics

Catch invalid-write panics while polling async KVStore futures so
the shared store helper still accepts debug assertions from stores
that validate inside the returned future.

Co-Authored-By: HAL 9000

Author: tnull

src/io/test_utils.rs

@@ -5,6 +5,7 @@
 // http://opensource.org/licenses/MIT>, at your option. You may not use this file except in
 // accordance with one or both of these licenses.
 
+use std::future::Future;
 use std::panic::RefUnwindSafe;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -168,6 +169,31 @@ pub(crate) fn random_storage_path() -> PathBuf {
 	temp_path
 }
 
+async fn catch_future_unwind<F: Future>(future: F) -> std::thread::Result<F::Output> {
+	let mut future = std::pin::pin!(future);
+	std::future::poll_fn(|cx| {
+		match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| future.as_mut().poll(cx))) {
+			Ok(std::task::Poll::Ready(output)) => std::task::Poll::Ready(Ok(output)),
+			Ok(std::task::Poll::Pending) => std::task::Poll::Pending,
+			Err(panic) => std::task::Poll::Ready(Err(panic)),
+		}
+	})
+	.await
+}
+
+async fn assert_invalid_write_fails<K: KVStore + RefUnwindSafe>(
+	kv_store: &K, primary_namespace: &str, secondary_namespace: &str, key: &str, data: Vec<u8>,
+) {
+	let res = std::panic::catch_unwind(|| {
+		KVStore::write(kv_store, primary_namespace, secondary_namespace, key, data)
+	});
+	if let Ok(fut) = res {
+		if let Ok(write_res) = catch_future_unwind(fut).await {
+			assert!(write_res.is_err());
+		}
+	}
+}
+
 pub(crate) async fn do_read_write_remove_list_persist<K: KVStore + RefUnwindSafe>(kv_store: &K) {
 	let data = vec![42u8; 32];
 
@@ -183,18 +209,9 @@ pub(crate) async fn do_read_write_remove_list_persist<K: KVStore + RefUnwindSafe
 	// Test empty primary/secondary namespaces are allowed, but not empty primary namespace and non-empty
 	// secondary primary_namespace, and not empty key.
 	KVStore::write(kv_store, "", "", key, data.clone()).await.unwrap();
-	let res = std::panic::catch_unwind(|| {
-		KVStore::write(kv_store, "", secondary_namespace, key, data.clone())
-	});
-	if let Ok(fut) = res {
-		assert!(fut.await.is_err());
-	}
-	let res = std::panic::catch_unwind(|| {
-		KVStore::write(kv_store, primary_namespace, secondary_namespace, "", data.clone())
-	});
-	if let Ok(fut) = res {
-		assert!(fut.await.is_err());
-	}
+	assert_invalid_write_fails(kv_store, "", secondary_namespace, key, data.clone()).await;
+	assert_invalid_write_fails(kv_store, primary_namespace, secondary_namespace, "", data.clone())
+		.await;
 
 	let listed_keys =
 		KVStore::list(kv_store, primary_namespace, secondary_namespace).await.unwrap();