code release

tests/test_shared_queue.py

@@ -0,0 +1,67 @@
+import sys
+import os
+
+ROOT_DIR = os.path.dirname(os.path.dirname(__file__))
+sys.path.append(ROOT_DIR)
+os.chdir(ROOT_DIR)
+
+import numpy as np
+from multiprocessing.managers import SharedMemoryManager
+from diffusion_policy.shared_memory.shared_memory_queue import SharedMemoryQueue, Full, Empty
+
+
+def test():
+    shm_manager = SharedMemoryManager()
+    shm_manager.start()
+    example = {
+        'cmd': 0,
+        'pose': np.zeros((6,))
+    }
+    queue = SharedMemoryQueue.create_from_examples(
+        shm_manager=shm_manager,
+        examples=example,
+        buffer_size=3
+    )
+    raised = False
+    try:
+        queue.get()
+    except Empty:
+        raised = True
+    assert raised
+
+    data = {
+        'cmd': 1,
+        'pose': np.ones((6,))
+    }
+    queue.put(data)
+    result = queue.get()
+    assert result['cmd'] == data['cmd']
+    assert np.allclose(result['pose'], data['pose'])
+
+    queue.put(data)
+    queue.put(data)
+    queue.put(data)
+    assert queue.qsize() == 3
+    raised = False
+    try:
+        queue.put(data)
+    except Full:
+        raised = True
+    assert raised
+
+    result = queue.get_all()
+    assert np.allclose(result['cmd'], [1,1,1])
+    
+    queue.put({'cmd': 0})
+    queue.put({'cmd': 1})
+    queue.put({'cmd': 2})
+    queue.get()
+    queue.put({'cmd': 3})
+
+    result = queue.get_k(3)
+    assert np.allclose(result['cmd'], [1,2,3])
+
+    queue.clear()
+
+if __name__ == "__main__":
+    test()