code: def test_heapset():
      heap = HeapSet(key=operator.attrgetter("i"))
  
      cx = C("x", 2)
      cy = C("y", 1)
      cz = C("z", 3)
      cw = C("w", 4)
      heap.add(cx)
      heap.add(cy)
      heap.add(cz)
      heap.add(cw)
      heap.add(C("x", 0))  # Ignored; x already in heap
      assert len(heap) == 4
      assert repr(heap) == "<HeapSet: 4 items>"
  
      assert cx in heap
      assert cy in heap
      assert cz in heap
      assert cw in heap
  
      heap_sorted = heap.sorted()
      # iteration does not empty heap
      assert len(heap) == 4
      assert next(heap_sorted) is cy
      assert next(heap_sorted) is cx
      assert next(heap_sorted) is cz
      assert next(heap_sorted) is cw
      with pytest.raises(StopIteration):
          next(heap_sorted)
  
      assert set(heap) == {cx, cy, cz, cw}
  
      assert heap.peek() is cy
      assert heap.pop() is cy
      assert cx in heap
      assert cy not in heap
      assert cz in heap
      assert cw in heap
  
      assert heap.peek() is cx
      assert heap.pop() is cx
      assert heap.pop() is cz
      assert heap.pop() is cw
      assert not heap
      with pytest.raises(KeyError):
          heap.pop()
      with pytest.raises(KeyError):
          heap.peek()
  
      # Test out-of-order discard
      heap.add(cx)
      heap.add(cy)
      heap.add(cz)
      heap.add(cw)
      assert heap.peek() is cy
  
      heap.remove(cy)
      assert cy not in heap
      with pytest.raises(KeyError):
          heap.remove(cy)
  
      heap.discard(cw)
      assert cw not in heap
      heap.discard(cw)
  
      assert len(heap) == 2
      assert list(heap.sorted()) == [cx, cz]
      # cy is at the top of heap._heap, but is skipped
      assert heap.peek() is cx
      assert heap.pop() is cx
      assert heap.peek() is cz
      assert heap.pop() is cz
      # heap._heap is not empty
      assert not heap
      with pytest.raises(KeyError):
          heap.peek()
      with pytest.raises(KeyError):
          heap.pop()
      assert list(heap.sorted()) == []
  
      # Test clear()
      heap.add(cx)
      heap.clear()
      assert not heap
      heap.add(cx)
      assert cx in heap
      # Test discard last element
      heap.discard(cx)
      assert not heap
      heap.add(cx)
      assert cx in heap
  
      # Test peekn()
      heap.add(cy)
      heap.add(cw)
      heap.add(cz)
      heap.add(cx)
      assert list(heap.peekn(3)) == [cy, cx, cz]
      heap.remove(cz)
      assert list(heap.peekn(10)) == [cy, cx, cw]
      assert list(heap.peekn(0)) == []
      assert list(heap.peekn(-1)) == []
      heap.remove(cy)
      assert list(heap.peekn(1)) == [cx]
      heap.remove(cw)
      assert list(heap.peekn(1)) == [cx]
      heap.remove(cx)
      assert list(heap.peekn(-1)) == []
      assert list(heap.peekn(0)) == []
      assert list(heap.peekn(1)) == []
      assert list(heap.peekn(2)) == []
  
      # Test resilience to failure in key()
      heap.add(cx)
      bad_key = C("bad_key", 0)
      del bad_key.i
      with pytest.raises(AttributeError):
          heap.add(bad_key)
      assert len(heap) == 1
      assert set(heap) == {cx}
  
      # Test resilience to failure in weakref.ref()
      class D:
          __slots__ = ("i",)
  
          def __init__(self, i):
              self.i = i
  
      with pytest.raises(TypeError):
          heap.add(D("bad_weakref", 2))
      assert len(heap) == 1
      assert set(heap) == {cx}
  
      # Test resilience to key() returning non-sortable output
      with pytest.raises(TypeError):
          heap.add(C("unsortable_key", None))
      assert len(heap) == 1
      assert set(heap) == {cx}
  
imports: from __future__ import annotations
  import heapq
  import operator
  import pickle
  import random
  from collections.abc import Mapping
  import pytest
  from distributed.collections import LRU, HeapSet, sum_mappings
label_desc: Use image node_id 2 for calling a global function with example usage: test_heapset() without return types