From d7c754ce49c8075596183ec199f7dac850e63140 Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Tue, 25 Mar 2025 00:10:48 +0800
Subject: [PATCH] failing test for UOp buffer ref count (#9563)

* failing test for UOp buffer ref count

* lint
---
 test/test_gc.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/test/test_gc.py b/test/test_gc.py
index cf90dc6201..a51c7adf18 100644
--- a/test/test_gc.py
+++ b/test/test_gc.py
@@ -79,5 +79,22 @@ class TestGC(unittest.TestCase):
       print(inspect.getclosurevars(UOp.toposort.fget))
       raise AssertionError(f"never gced {[x for x in gc.get_objects() if isinstance(x, Buffer)]}")
 
+  def test_buffer_refcount(self):
+    init = bufs_allocated()
+    a = Tensor.empty(10)
+    self.assertEqual(bufs_allocated()-init, 0)
+    a.realize()
+    real_buf = a.lazydata.buffer
+    # after the Tensor UOp is deleted there shouldn't be any references on the Buffer
+    with self.assertRaises(AssertionError):
+      self.assertEqual(real_buf.lb_refcount, 1)
+    self.assertEqual(bufs_allocated()-init, 1)
+    del a.lazydata
+    with self.assertRaises(AssertionError):
+      self.assertEqual(real_buf.lb_refcount, 0)
+    self.assertEqual(bufs_allocated()-init, 1) # keep the buffer alive
+    del real_buf
+    self.assertEqual(bufs_allocated()-init, 0)
+
 if __name__ == '__main__':
   unittest.main()