4 files changed, 87 insertions, 4 deletions

diff --git a/lib/Utils/tests/.runsettings b/lib/Utils/tests/.runsettings
new file mode 100644
index 0000000..a4d0377
--- /dev/null
+++ b/lib/Utils/tests/.runsettings
@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="utf-8"?>
+<RunSettings>
+    <RunConfiguration>
+        <EnvironmentVariables>
+            <VNLIB_SHARED_HEAP_DIAGNOSTICS>1</VNLIB_SHARED_HEAP_DIAGNOSTICS>
+        </EnvironmentVariables>
+    </RunConfiguration>
+</RunSettings>
\ No newline at end of file
diff --git a/lib/Utils/tests/Memory/MemoryHandleTest.cs b/lib/Utils/tests/Memory/MemoryHandleTest.cs
index 34dbb60..2ef8a41 100644
--- a/lib/Utils/tests/Memory/MemoryHandleTest.cs
+++ b/lib/Utils/tests/Memory/MemoryHandleTest.cs
@@ -42,7 +42,7 @@ namespace VNLib.Utils.Memory.Tests
             Assert.ThrowsException<ArgumentOutOfRangeException>(() => Shared.Alloc<byte>(-1));
 
             //Make sure over-alloc throws
-            Assert.ThrowsException<NativeMemoryOutOfMemoryException>(() => Shared.Alloc<byte>(nuint.MaxValue, false));
+            Assert.ThrowsException<OutOfMemoryException>(() => Shared.Alloc<byte>(nuint.MaxValue, false));
         }
 #if TARGET_64_BIT
         [TestMethod]
diff --git a/lib/Utils/tests/Memory/MemoryUtilTests.cs b/lib/Utils/tests/Memory/MemoryUtilTests.cs
index fb3700e..d55817c 100644
--- a/lib/Utils/tests/Memory/MemoryUtilTests.cs
+++ b/lib/Utils/tests/Memory/MemoryUtilTests.cs
@@ -5,6 +5,7 @@ using System.Security.Cryptography;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 
 using VNLib.Utils.Extensions;
+using VNLib.Utils.Memory.Diagnostics;
 
 namespace VNLib.Utils.Memory.Tests
 {
@@ -243,7 +244,7 @@ namespace VNLib.Utils.Memory.Tests
                 //Test pinned outsie handle size
                 Assert.ThrowsException<ArgumentOutOfRangeException>(() => _ = handle.Pin(1024));
             }
-           
+
 
             //Negative value
             Assert.ThrowsException<ArgumentException>(() => _ = MemoryUtil.SafeAlloc<byte>(-1));
@@ -329,5 +330,79 @@ namespace VNLib.Utils.Memory.Tests
             //Free struct
             MemoryUtil.Shared.StructFree(s);
         }
+
+        [TestMethod()]
+        public void GetSharedHeapStatsTest()
+        {
+            //Confirm heap diagnostics are enabled
+            Assert.AreEqual<string?>(Environment.GetEnvironmentVariable(MemoryUtil.SHARED_HEAP_ENABLE_DIAGNOISTICS_ENV), "1");
+
+            //Get current stats
+            HeapStatistics preTest = MemoryUtil.GetSharedHeapStats();
+
+            //Alloc block
+            using IMemoryHandle<byte> handle = MemoryUtil.Shared.Alloc<byte>(1024);
+
+            //Get stats
+            HeapStatistics postTest = MemoryUtil.GetSharedHeapStats();
+
+            Assert.IsTrue(postTest.AllocatedBytes == preTest.AllocatedBytes + 1024);
+            Assert.IsTrue(postTest.AllocatedBlocks == preTest.AllocatedBlocks + 1);
+
+            //Free block
+            handle.Dispose();
+
+            //Get stats
+            HeapStatistics postFree = MemoryUtil.GetSharedHeapStats();
+
+            //Confirm stats are back to pre test
+            Assert.IsTrue(preTest.AllocatedBytes == postFree.AllocatedBytes);
+            Assert.IsTrue(preTest.AllocatedBlocks == postFree.AllocatedBlocks);
+        }
+
+        [TestMethod()]
+        public void DiagnosticsHeapWraperTest()
+        {
+            //Get a fresh heap
+            IUnmangedHeap heap = MemoryUtil.InitializeNewHeapForProcess();
+
+            //Init wrapper and dispose
+            using TrackedHeapWrapper wrapper = new(heap);
+
+            //Confirm 0 stats
+            HeapStatistics preTest = wrapper.GetCurrentStats();
+
+            Assert.IsTrue(preTest.AllocatedBytes == 0);
+            Assert.IsTrue(preTest.AllocatedBlocks == 0);
+            Assert.IsTrue(preTest.MaxHeapSize == 0);
+            Assert.IsTrue(preTest.MaxBlockSize == 0);
+            Assert.IsTrue(preTest.MinBlockSize == ulong.MaxValue);
+
+            //Alloc a test block
+            using IMemoryHandle<byte> handle = wrapper.Alloc<byte>(1024);
+
+            //Get stats
+            HeapStatistics postTest = wrapper.GetCurrentStats();
+
+            //Confirm stats represent a single block
+            Assert.IsTrue(postTest.AllocatedBytes == 1024);
+            Assert.IsTrue(postTest.AllocatedBlocks == 1);
+            Assert.IsTrue(postTest.MaxHeapSize == 1024);
+            Assert.IsTrue(postTest.MaxBlockSize == 1024);
+            Assert.IsTrue(postTest.MinBlockSize == 1024);
+
+            //Free the block
+            handle.Dispose();
+
+            //Get stats
+            HeapStatistics postFree = wrapper.GetCurrentStats();
+
+            //Confirm stats are back to 0, or represent the single block
+            Assert.IsTrue(postFree.AllocatedBytes == 0);
+            Assert.IsTrue(postFree.AllocatedBlocks == 0);
+            Assert.IsTrue(postFree.MaxHeapSize == 1024);
+            Assert.IsTrue(postFree.MaxBlockSize == 1024);
+            Assert.IsTrue(postFree.MinBlockSize == 1024);
+        }
     }
 }
\ No newline at end of file
diff --git a/lib/Utils/tests/Memory/VnTableTests.cs b/lib/Utils/tests/Memory/VnTableTests.cs
index c9f99ea..6f9fbf2 100644
--- a/lib/Utils/tests/Memory/VnTableTests.cs
+++ b/lib/Utils/tests/Memory/VnTableTests.cs
@@ -53,9 +53,9 @@ namespace VNLib.Utils.Memory.Tests
 
             
             //Test oom, should be native
-            Assert.ThrowsException<NativeMemoryOutOfMemoryException>(() =>
+            Assert.ThrowsException<OutOfMemoryException>(() =>
             {
-                using VnTable<int> table = new(int.MaxValue, 2);
+                using VnTable<int> table = new(uint.MaxValue, 3);
             });
         }