From 5413e30f572ebf06f27a5ad198d02444eed94da6 Mon Sep 17 00:00:00 2001
From: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>
Date: Sun, 1 Feb 2026 17:30:01 -0500
Subject: [PATCH] tests: harden memory leak check in test_siblings_without_eval

When running the test suite on darwin as part of nixpkgs package build,
we often hit the following error:

```
       >         mx.synchronize()
       >         t()
       >         gc.collect()
       >         expected = get_mem()
       >         for _ in range(100):
       >             t()
       >         used = get_mem()
       > >       self.assertEqual(expected, used)
       > E       AssertionError: 273678336 != 273694720
```

You can find an example of the full failure log here:
https://hydra.nixos.org/build/320731525

The difference between actual and expected is 16k (one page size on
aarch64-darwin).

The test seems brittle as it depends on *peak* memory, not real memory
at the moment of measurement, so anything during the run of the process
that could push the allocator to allocate another page can break the
test.

This patch switches the test case to measure real memory using psutil on
all platforms.

Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>
---
 python/tests/test_array.py | 14 +++-----------
 1 file changed, 3 insertions(+), 11 deletions(-)

diff --git a/python/tests/test_array.py b/python/tests/test_array.py
index 4efed9dac9..12ed0d6f2d 100644
--- a/python/tests/test_array.py
+++ b/python/tests/test_array.py
@@ -4,21 +4,16 @@
 import operator
 import os
 import pickle
-import platform
 import sys
 import unittest
 import weakref
 from copy import copy, deepcopy
 from itertools import permutations
 
-if platform.system() == "Windows":
-    import psutil
-else:
-    import resource
-
 import mlx.core as mx
 import mlx_tests
 import numpy as np
+import psutil
 
 try:
     import tensorflow as tf
@@ -2087,11 +2082,8 @@ def test_deep_graphs(self):
 
     def test_siblings_without_eval(self):
         def get_mem():
-            if platform.system() == "Windows":
-                process = psutil.Process(os.getpid())
-                return process.memory_info().peak_wset
-            else:
-                return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+            process = psutil.Process(os.getpid())
+            return process.memory_info().rss
 
         key = mx.array([1, 2])