def test_output_funnel_drain(funnel: OutputFunnel) -> None:
with funnel.plugged():
funnel.write(b"A")
text = funnel.drain()
assert text == "A"
funnel.write(b"B")
assert response_texts(funnel) == [["B"]]
assert written(funnel) == b"B"
def test_output_funnel_context_drain(funnel: OutputFunnel) -> None:
funnel.write(b"A")
assert written(funnel) == b"A"
with funnel.plugged():
funnel.write(b"B")
assert response_texts(funnel) == [["B"]]
code = funnel.drain()
assert response_texts(funnel) == [[]]
assert code == "B"
assert written(funnel) == b"A"