Add tests to check the size of memory allocations

src/memory/device/device_buffer.rs

@@ -459,4 +459,22 @@ mod test_device_buffer {
             let _slice = &buffer[0..5];
         }
     }
+
+    #[test]
+    fn test_allocate_correct_size() {
+        use crate::context::CurrentContext;
+
+        let _context = crate::quick_init().unwrap();
+        let total_memory = CurrentContext::get_device()
+            .unwrap()
+            .total_memory()
+            .unwrap();
+
+        // Don't allocate all memory to leave some space for the display's frame buffer
+        let allocation_size = (total_memory * 3) / 4 / mem::size_of::<u64>();
+        unsafe {
+            // Test if allocation fails with an out-of-memory error
+            let _buffer = DeviceBuffer::<u64>::uninitialized(allocation_size).unwrap();
+        };
+    }
 }

src/memory/locked.rs

@@ -332,4 +332,16 @@ mod test {
         let err = LockedBuffer::new(&0u64, ::std::usize::MAX - 1).unwrap_err();
         assert_eq!(CudaError::InvalidMemoryAllocation, err);
     }
+
+    #[test]
+    fn test_allocate_correct_size() {
+        let _context = crate::quick_init().unwrap();
+
+        // Placeholder - read out available system memory here
+        let allocation_size = 1;
+        unsafe {
+            // Test if allocation fails with an out-of-memory error
+            let _buffer = LockedBuffer::<u64>::uninitialized(allocation_size).unwrap();
+        }
+    }
 }

src/memory/unified.rs

@@ -719,4 +719,22 @@ mod test_unified_buffer {
         let err = UnifiedBuffer::new(&0u64, ::std::usize::MAX - 1).unwrap_err();
         assert_eq!(CudaError::InvalidMemoryAllocation, err);
     }
+
+    #[test]
+    fn test_allocate_correct_size() {
+        use crate::context::CurrentContext;
+
+        let _context = crate::quick_init().unwrap();
+        let total_memory = CurrentContext::get_device()
+            .unwrap()
+            .total_memory()
+            .unwrap();
+
+        // Don't allocate all memory to leave some space for the display's frame buffer
+        let allocation_size = (total_memory * 3) / 4 / mem::size_of::<u64>();
+        unsafe {
+            // Test if allocation fails with an out-of-memory error
+            let _buffer = UnifiedBuffer::<u64>::uninitialized(allocation_size).unwrap();
+        }
+    }
 }