Cleanup Custom Tests

This moves the tests for the custom RNG registration into custom.rs. It
also:
  - Simplifies the registered custom RNG
  - Makes sure the custom RNG _is not_ used on supported platforms
  - Makes sure the custom RNG _is_ used on unsupported platforms

Signed-off-by: Joe Richey <[email protected]>

Author: josephlr

src/custom.rs

@@ -103,3 +103,39 @@ pub fn getrandom_inner(dest: &mut [MaybeUninit<u8>]) -> Result<(), Error> {
         Some(code) => Err(Error::from(code)),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const LEN7_CODE: u32 = Error::CUSTOM_START + 7;
+    // Returns a custom error if input is length 7, otherwise fills with 0x55.
+    fn insecure_rng(dest: &mut [u8]) -> Result<(), Error> {
+        if dest.len() == 7 {
+            return Err(NonZeroU32::new(LEN7_CODE).unwrap().into());
+        }
+        dest.fill(0x55);
+        Ok(())
+    }
+
+    // Test registering a custom implementation, even on supported platforms.
+    register_custom_getrandom!(insecure_rng);
+
+    // On a supported platform, make sure the custom implementation isn't used.
+    #[cfg(target_os = "linux")]
+    #[test]
+    fn custom_not_used() {
+        crate::getrandom(&mut [0; 7]).unwrap();
+    }
+    // On an unsupported platform, make sure the custom implementation is used.
+    #[cfg(all(target_arch = "wasm32", target_os = "unknown", not(feature = "js")))]
+    #[wasm_bindgen_test::wasm_bindgen_test]
+    fn custom_used() {
+        let err = crate::getrandom(&mut [0; 7]).unwrap_err();
+        assert_eq!(err.code().get(), LEN7_CODE);
+
+        let mut buf = [0; 12];
+        crate::getrandom(&mut buf).unwrap();
+        assert_eq!(buf, [0x55; 12]);
+    }
+}