1  import re
   2  from typing import Tuple
   3  
   4  import pytest
   5  
   6  from .._receivebuffer import ReceiveBuffer
   7  
   8  
   9  def test_receivebuffer() -> None:
  10      b = ReceiveBuffer()
  11      assert not b
  12      assert len(b) == 0
  13      assert bytes(b) == b""
  14  
  15      b += b"123"
  16      assert b
  17      assert len(b) == 3
  18      assert bytes(b) == b"123"
  19  
  20      assert bytes(b) == b"123"
  21  
  22      assert b.maybe_extract_at_most(2) == b"12"
  23      assert b
  24      assert len(b) == 1
  25      assert bytes(b) == b"3"
  26  
  27      assert bytes(b) == b"3"
  28  
  29      assert b.maybe_extract_at_most(10) == b"3"
  30      assert bytes(b) == b""
  31  
  32      assert b.maybe_extract_at_most(10) is None
  33      assert not b
  34  
  35      ################################################################
  36      # maybe_extract_until_next
  37      ################################################################
  38  
  39      b += b"123\n456\r\n789\r\n"
  40  
  41      assert b.maybe_extract_next_line() == b"123\n456\r\n"
  42      assert bytes(b) == b"789\r\n"
  43  
  44      assert b.maybe_extract_next_line() == b"789\r\n"
  45      assert bytes(b) == b""
  46  
  47      b += b"12\r"
  48      assert b.maybe_extract_next_line() is None
  49      assert bytes(b) == b"12\r"
  50  
  51      b += b"345\n\r"
  52      assert b.maybe_extract_next_line() is None
  53      assert bytes(b) == b"12\r345\n\r"
  54  
  55      # here we stopped at the middle of b"\r\n" delimiter
  56  
  57      b += b"\n6789aaa123\r\n"
  58      assert b.maybe_extract_next_line() == b"12\r345\n\r\n"
  59      assert b.maybe_extract_next_line() == b"6789aaa123\r\n"
  60      assert b.maybe_extract_next_line() is None
  61      assert bytes(b) == b""
  62  
  63      ################################################################
  64      # maybe_extract_lines
  65      ################################################################
  66  
  67      b += b"123\r\na: b\r\nfoo:bar\r\n\r\ntrailing"
  68      lines = b.maybe_extract_lines()
  69      assert lines == [b"123", b"a: b", b"foo:bar"]
  70      assert bytes(b) == b"trailing"
  71  
  72      assert b.maybe_extract_lines() is None
  73  
  74      b += b"\r\n\r"
  75      assert b.maybe_extract_lines() is None
  76  
  77      assert b.maybe_extract_at_most(100) == b"trailing\r\n\r"
  78      assert not b
  79  
  80      # Empty body case (as happens at the end of chunked encoding if there are
  81      # no trailing headers, e.g.)
  82      b += b"\r\ntrailing"
  83      assert b.maybe_extract_lines() == []
  84      assert bytes(b) == b"trailing"
  85  
  86  
  87  @pytest.mark.parametrize(
  88      "data",
  89      [
  90          pytest.param(
  91              (
  92                  b"HTTP/1.1 200 OK\r\n",
  93                  b"Content-type: text/plain\r\n",
  94                  b"Connection: close\r\n",
  95                  b"\r\n",
  96                  b"Some body",
  97              ),
  98              id="with_crlf_delimiter",
  99          ),
 100          pytest.param(
 101              (
 102                  b"HTTP/1.1 200 OK\n",
 103                  b"Content-type: text/plain\n",
 104                  b"Connection: close\n",
 105                  b"\n",
 106                  b"Some body",
 107              ),
 108              id="with_lf_only_delimiter",
 109          ),
 110          pytest.param(
 111              (
 112                  b"HTTP/1.1 200 OK\n",
 113                  b"Content-type: text/plain\r\n",
 114                  b"Connection: close\n",
 115                  b"\n",
 116                  b"Some body",
 117              ),
 118              id="with_mixed_crlf_and_lf",
 119          ),
 120      ],
 121  )
 122  def test_receivebuffer_for_invalid_delimiter(data: Tuple[bytes]) -> None:
 123      b = ReceiveBuffer()
 124  
 125      for line in data:
 126          b += line
 127  
 128      lines = b.maybe_extract_lines()
 129  
 130      assert lines == [
 131          b"HTTP/1.1 200 OK",
 132          b"Content-type: text/plain",
 133          b"Connection: close",
 134      ]
 135      assert bytes(b) == b"Some body"
 136