From c399e3163daa8a74e7063b7e77639969432e52c5 Mon Sep 17 00:00:00 2001
From: Shrey Amin <samin@cloudflare.com>
Date: Sun, 3 Nov 2024 20:51:38 -0500
Subject: [PATCH] Add test for disk I/O usage

---
 Cargo.toml         |  1 +
 tests/disk.rs      | 79 ++++++++++++++++++++++++++++++++++++++++++++++
 tests/disk_list.rs | 16 ----------
 3 files changed, 80 insertions(+), 16 deletions(-)
 create mode 100644 tests/disk.rs
 delete mode 100644 tests/disk_list.rs

diff --git a/Cargo.toml b/Cargo.toml
index 94f18f68d..806bfe520 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -117,6 +117,7 @@ procfs = "0.17.0"
 [dev-dependencies]
 serde_json = "1.0" # Used in documentation tests.
 bstr = "1.9.0"
+tempfile = "3.9"
 
 [[example]]
 name = "simple"
diff --git a/tests/disk.rs b/tests/disk.rs
new file mode 100644
index 000000000..63463c587
--- /dev/null
+++ b/tests/disk.rs
@@ -0,0 +1,79 @@
+// Take a look at the license at the top of the repository in the LICENSE file.
+
+#[test]
+#[cfg(all(feature = "system", feature = "disk"))]
+fn test_disks() {
+    if sysinfo::IS_SUPPORTED_SYSTEM {
+        let s = sysinfo::System::new_all();
+        // If we don't have any physical core present, it's very likely that we're inside a VM...
+        if s.physical_core_count().unwrap_or_default() > 0 {
+            let mut disks = sysinfo::Disks::new();
+            assert!(disks.list().is_empty());
+            disks.refresh_list();
+            assert!(!disks.list().is_empty());
+        }
+    }
+}
+
+#[test]
+#[cfg(feature = "disk")]
+fn test_disks_usage() {
+    use std::io::Write;
+    use tempfile::NamedTempFile;
+
+    let s = sysinfo::System::new_all();
+    dbg!(s.physical_core_count());
+
+    if !sysinfo::IS_SUPPORTED_SYSTEM || s.physical_core_count().unwrap_or_default() == 0 {
+        return;
+    }
+
+    let mut disks = sysinfo::Disks::new_with_refreshed_list();
+    for disk in disks.list() {
+        println!("BEFORE: {disk:?}");
+    }
+
+    let diskstats = std::fs::read_to_string("/proc/diskstats").unwrap();
+    println!("PROCSTATS BEFORE:\n{}", diskstats);
+
+    let temp_dir = tempfile::tempdir().unwrap();
+
+    let mut file = std::fs::File::create(temp_dir.path().join("test.txt")).unwrap();
+
+    // Write 10mb worth of data to the temp file. Repeat this 10 times to increase the chances
+    // of the OS registering the disk writes.
+    for _ in 0..10 {
+        let data = vec![1u8; 10 * 1024 * 1024];
+        file.write_all(&data).unwrap();
+        // The sync_all call is important to ensure all the data is persisted to disk. Without
+        // the call, this test is flaky.
+        file.sync_all().unwrap();
+    }
+
+    disks.refresh();
+
+    let diskstats_after = std::fs::read_to_string("/proc/diskstats").unwrap();
+
+    println!("PROCSTATS AFTER:\n{}", diskstats_after);
+
+    for disk in disks.list() {
+        println!("AFTER: {disk:?}");
+    }
+    // Depending on the OS and how disks are configured, the disk usage may be the exact same
+    // across multiple disks. To account for this, collect the disk usages and dedup
+    let mut disk_usages = disks.list().iter().map(|d| d.usage()).collect::<Vec<_>>();
+    disk_usages.dedup();
+
+    let mut written_bytes = 0;
+    for disk_usage in disk_usages {
+        written_bytes += disk_usage.written_bytes;
+    }
+
+    // written_bytes should have increased by about 10mb, but this is not fully reliable in CI Linux. For now,
+    // just verify the number is non-zero.
+    #[cfg(not(target_os = "freebsd"))]
+    assert!(written_bytes > 0);
+    // Disk usage is not yet supported on freebsd
+    #[cfg(target_os = "freebsd")]
+    assert_eq!(written_bytes, 0);
+}
diff --git a/tests/disk_list.rs b/tests/disk_list.rs
deleted file mode 100644
index 9ddaefb50..000000000
--- a/tests/disk_list.rs
+++ /dev/null
@@ -1,16 +0,0 @@
-// Take a look at the license at the top of the repository in the LICENSE file.
-
-#[test]
-#[cfg(all(feature = "system", feature = "disk"))]
-fn test_disks() {
-    if sysinfo::IS_SUPPORTED_SYSTEM {
-        let s = sysinfo::System::new_all();
-        // If we don't have any physical core present, it's very likely that we're inside a VM...
-        if s.physical_core_count().unwrap_or_default() > 0 {
-            let mut disks = sysinfo::Disks::new();
-            assert!(disks.list().is_empty());
-            disks.refresh_list();
-            assert!(!disks.list().is_empty());
-        }
-    }
-}