forked from goodboy/tractor
1
0
Fork 0

Add some bcaster ref sanity asserts around subscriptions

live_on_air_from_tokio
Tyler Goodlet 2021-08-31 13:06:17 -04:00
parent 0d70e3081a
commit 63ec740e27
1 changed files with 4 additions and 2 deletions

View File

@ -9,7 +9,7 @@ from dataclasses import dataclass
from typing import ( from typing import (
Any, Iterator, Optional, Callable, Any, Iterator, Optional, Callable,
AsyncGenerator, Dict, AsyncGenerator, Dict,
AsyncIterator, Awaitable AsyncIterator
) )
import warnings import warnings
@ -264,7 +264,7 @@ class ReceiveMsgStream(trio.abc.ReceiveChannel):
self, self,
# use memory channel size by default # use memory channel size by default
self._rx_chan._state.max_buffer_size, # type: ignore self._rx_chan._state.max_buffer_size, # type: ignore
receive_afunc=self.receive, receive_afunc=self.receive,
) )
# NOTE: we override the original stream instance's receive # NOTE: we override the original stream instance's receive
@ -277,6 +277,8 @@ class ReceiveMsgStream(trio.abc.ReceiveChannel):
# https://github.com/python/mypy/issues/708 # https://github.com/python/mypy/issues/708
async with self._broadcaster.subscribe() as bstream: async with self._broadcaster.subscribe() as bstream:
assert bstream.key != self._broadcaster.key
assert bstream._recv == self._broadcaster._recv
yield bstream yield bstream