68 lines
1.4 KiB
Python
68 lines
1.4 KiB
Python
|
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()
|