Merge pull request #77 from kaznak/202107.kaznak.store_wait-with-data

The wait method returns the stream data (#75)

pybotters/store.py

@@ -25,7 +25,7 @@ def __init__(self, keys: List[str] = [], data: List[Item] = []) -> None:
         self._data: Dict[uuid.UUID, Item] = {}
         self._index: Dict[int, uuid.UUID] = {}
         self._keys: Tuple[str, ...] = tuple(keys if keys else self._KEYS)
-        self._events: List[asyncio.Event] = []
+        self._events: Dict[asyncio.Event, List[Item]] = {}
         self._insert(data)
         if hasattr(self, '_init'):
             getattr(self, '_init')()
@@ -61,7 +61,8 @@ def _insert(self, data: List[Item]) -> None:
                 _id = uuid.uuid4()
                 self._data[_id] = item
             self._sweep_without_key()
-        self._set()
+        # !TODO! This behaviour might be undesirable.
+        self._set(data)
 
     def _update(self, data: List[Item]) -> None:
         if self._keys:
@@ -84,7 +85,8 @@ def _update(self, data: List[Item]) -> None:
                 _id = uuid.uuid4()
                 self._data[_id] = item
             self._sweep_without_key()
-        self._set()
+        # !TODO! This behaviour might be undesirable.
+        self._set(data)
 
     def _delete(self, data: List[Item]) -> None:
         if self._keys:
@@ -98,7 +100,8 @@ def _delete(self, data: List[Item]) -> None:
                     if keyhash in self._index:
                         del self._data[self._index[keyhash]]
                         del self._index[keyhash]
-        self._set()
+        # !TODO! This behaviour might be undesirable.
+        self._set(data)
 
     def _clear(self) -> None:
         self._data.clear()
@@ -143,15 +146,18 @@ def find(self, query: Item = {}) -> List[Item]:
         else:
             return list(self)
 
-    def _set(self) -> None:
+    def _set(self, data: List[Item] = None) -> None:
         for event in self._events:
             event.set()
-        self._events.clear()
+            self._events[event].extend(data)
 
-    async def wait(self) -> None:
+    async def wait(self) -> List[Item]:
         event = asyncio.Event()
-        self._events.append(event)
+        ret = []
+        self._events[event] = ret
         await event.wait()
+        del self._events[event]
+        return ret
 
 
 TDataStore = TypeVar('TDataStore', bound=DataStore)

tests/test_store.py

@@ -212,21 +212,89 @@ def test__iter__():
 
 def test_set():
     ds = pybotters.store.DataStore()
-    events = [asyncio.Event(), asyncio.Event(), asyncio.Event()]
-    ds._events.extend(events)
-    ds._set()
-    assert all(e.is_set() for e in events)
-    assert not len(ds._events)
+    event = asyncio.Event()
+    ds._events[event] = []
+    data = [{'dummy1': 'data1'}, {'dummy2': 'data2'}, {'dummy3': 'data3'}]
+    ds._set(data)
+    assert all(e.is_set() for e in ds._events)
+    assert ds._events[event] == data
 
 
 @pytest.mark.asyncio
-async def test_wait():
+async def test_wait_set():
+    data = [{'dummy': 'data'}]
+    ret = {}
+
     class DataStoreHasDummySet(pybotters.store.DataStore):
-        async def _set(self) -> None:
-            return super()._set()
-
-    ds = DataStoreHasDummySet()
-    t_wait = asyncio.create_task(ds.wait())
-    t_set = asyncio.create_task(ds._set())
-    await asyncio.wait_for(t_wait, timeout=5.0)
-    assert t_set.done()
+        async def _set(self, data) -> None:
+            return super()._set(data)
+
+    async def wait_func(ds):
+        ret['val'] = await ds.wait()
+
+    ds0 = DataStoreHasDummySet()
+    t_wait0 = asyncio.create_task(wait_func(ds0))
+    t_set0 = asyncio.create_task(ds0._set(data))
+    await asyncio.wait_for(t_wait0, timeout=5.0)
+    assert t_set0.done()
+    assert data == ret['val']
+
+
+@pytest.mark.asyncio
+async def test_wait_insert():
+    data = [{'dummy': 'data'}]
+    ret = {}
+
+    class DataStoreHasDummyInsert(pybotters.store.DataStore):
+        async def _insert(self, data) -> None:
+            return super()._insert(data)
+
+    async def wait_func(ds):
+        ret['val'] = await ds.wait()
+
+    ds1 = DataStoreHasDummyInsert()
+    t_wait1 = asyncio.create_task(wait_func(ds1))
+    t_set1 = asyncio.create_task(ds1._insert(data))
+    await asyncio.wait_for(t_wait1, timeout=5.0)
+    assert t_set1.done()
+    assert data == ret['val']
+
+
+@pytest.mark.asyncio
+async def test_wait_update():
+    data = [{'dummy': 'data'}]
+    ret = {}
+
+    class DataStoreHasDummyUpdate(pybotters.store.DataStore):
+        async def _update(self, data) -> None:
+            return super()._update(data)
+
+    async def wait_func(ds):
+        ret['val'] = await ds.wait()
+
+    ds2 = DataStoreHasDummyUpdate()
+    t_wait2 = asyncio.create_task(wait_func(ds2))
+    t_set2 = asyncio.create_task(ds2._update(data))
+    await asyncio.wait_for(t_wait2, timeout=5.0)
+    assert t_set2.done()
+    assert data == ret['val']
+
+
+@pytest.mark.asyncio
+async def test_wait_delete():
+    data = [{'dummy': 'data'}]
+    ret = {}
+
+    class DataStoreHasDummyDelete(pybotters.store.DataStore):
+        async def _delete(self, data) -> None:
+            return super()._delete(data)
+
+    async def wait_func(ds):
+        ret['val'] = await ds.wait()
+
+    ds3 = DataStoreHasDummyDelete()
+    t_wait3 = asyncio.create_task(wait_func(ds3))
+    t_set3 = asyncio.create_task(ds3._delete(data))
+    await asyncio.wait_for(t_wait3, timeout=5.0)
+    assert t_set3.done()
+    assert data == ret['val']