airflow_scheduler/.venv/Lib/site-packages/h11/tests/helpers.py

from typing import cast, List, Type, Union, ValuesView

from .._connection import Connection, NEED_DATA, PAUSED
from .._events import (
    ConnectionClosed,
    Data,
    EndOfMessage,
    Event,
    InformationalResponse,
    Request,
    Response,
)
from .._state import CLIENT, CLOSED, DONE, MUST_CLOSE, SERVER
from .._util import Sentinel

try:
    from typing import Literal
except ImportError:
    from typing_extensions import Literal  # type: ignore


def get_all_events(conn: Connection) -> List[Event]:
    got_events = []
    while True:
        event = conn.next_event()
        if event in (NEED_DATA, PAUSED):
            break
        event = cast(Event, event)
        got_events.append(event)
        if type(event) is ConnectionClosed:
            break
    return got_events


def receive_and_get(conn: Connection, data: bytes) -> List[Event]:
    conn.receive_data(data)
    return get_all_events(conn)


# Merges adjacent Data events, converts payloads to bytestrings, and removes
# chunk boundaries.
def normalize_data_events(in_events: List[Event]) -> List[Event]:
    out_events: List[Event] = []
    for event in in_events:
        if type(event) is Data:
            event = Data(data=bytes(event.data), chunk_start=False, chunk_end=False)
        if out_events and type(out_events[-1]) is type(event) is Data:
            out_events[-1] = Data(
                data=out_events[-1].data + event.data,
                chunk_start=out_events[-1].chunk_start,
                chunk_end=out_events[-1].chunk_end,
            )
        else:
            out_events.append(event)
    return out_events


# Given that we want to write tests that push some events through a Connection
# and check that its state updates appropriately... we might as make a habit
# of pushing them through two Connections with a fake network link in
# between.
class ConnectionPair:
    def __init__(self) -> None:
        self.conn = {CLIENT: Connection(CLIENT), SERVER: Connection(SERVER)}
        self.other = {CLIENT: SERVER, SERVER: CLIENT}

    @property
    def conns(self) -> ValuesView[Connection]:
        return self.conn.values()

    # expect="match" if expect=send_events; expect=[...] to say what expected
    def send(
        self,
        role: Type[Sentinel],
        send_events: Union[List[Event], Event],
        expect: Union[List[Event], Event, Literal["match"]] = "match",
    ) -> bytes:
        if not isinstance(send_events, list):
            send_events = [send_events]
        data = b""
        closed = False
        for send_event in send_events:
            new_data = self.conn[role].send(send_event)
            if new_data is None:
                closed = True
            else:
                data += new_data
        # send uses b"" to mean b"", and None to mean closed
        # receive uses b"" to mean closed, and None to mean "try again"
        # so we have to translate between the two conventions
        if data:
            self.conn[self.other[role]].receive_data(data)
        if closed:
            self.conn[self.other[role]].receive_data(b"")
        got_events = get_all_events(self.conn[self.other[role]])
        if expect == "match":
            expect = send_events
        if not isinstance(expect, list):
            expect = [expect]
        assert got_events == expect
        return data