123 lines
4.5 KiB
Python
123 lines
4.5 KiB
Python
import unittest
|
|
|
|
from source_conversion import (
|
|
InvalidSourceError,
|
|
SourceTooLargeError,
|
|
convert_source,
|
|
iter_completed_bytes,
|
|
)
|
|
|
|
|
|
class FakeUpstream:
|
|
def __init__(self, body, headers=None, fail_after=None):
|
|
self.body = body
|
|
self.headers = headers or {"Content-Type": "text/plain; charset=utf-8"}
|
|
self.fail_after = fail_after
|
|
|
|
def iter_content(self, chunk_size):
|
|
for i in range(0, len(self.body), chunk_size):
|
|
yield self.body[i:i + chunk_size]
|
|
if self.fail_after is not None and i >= self.fail_after:
|
|
raise IOError("upstream disconnected")
|
|
|
|
|
|
def convert_test_line(line, ip):
|
|
line = line.strip()
|
|
if not line or line.startswith('#'):
|
|
return None
|
|
if line.startswith('||') and line.endswith('^'):
|
|
return f'{ip} {line[2:-1]}'
|
|
if line.startswith('0.0.0.0 '):
|
|
return f'{ip} {line.split()[1]}'
|
|
if line.endswith('.test'):
|
|
return f'{ip} {line}'
|
|
return None
|
|
|
|
|
|
def run_conversion(body, *, headers=None, chunk=7, source_limit=1024, output_limit=1024,
|
|
line_limit=128, debug=None, fail_after=None):
|
|
return convert_source(
|
|
FakeUpstream(body, headers=headers, fail_after=fail_after),
|
|
"0.0.0.0",
|
|
convert_test_line,
|
|
read_chunk=chunk,
|
|
line_limit=line_limit,
|
|
source_limit=source_limit,
|
|
output_limit=output_limit,
|
|
spool_memory_limit=16,
|
|
debug_lines=debug,
|
|
)
|
|
|
|
|
|
class ConvertSourceTests(unittest.TestCase):
|
|
def test_valid_host_list_is_bytes_and_completed(self):
|
|
src = b'# comment\n0.0.0.0 first.test\n||second.test^\n'
|
|
result = run_conversion(src)
|
|
body_chunks = list(iter_completed_bytes(result.output, 5))
|
|
self.assertTrue(all(isinstance(v, bytes) for v in body_chunks))
|
|
self.assertEqual(b''.join(body_chunks), b'0.0.0.0 first.test\n0.0.0.0 second.test\n')
|
|
self.assertTrue(result.output.closed)
|
|
self.assertEqual(result.lines_emitted, 2)
|
|
|
|
def test_octet_stream_without_charset(self):
|
|
result = run_conversion(
|
|
b'example.test\n', headers={"Content-Type": "application/octet-stream"}
|
|
)
|
|
self.assertEqual(result.output.read(), b'0.0.0.0 example.test\n')
|
|
result.output.close()
|
|
|
|
def test_utf8_bom(self):
|
|
result = run_conversion(b'\xef\xbb\xbfexample.test\n')
|
|
self.assertEqual(result.output.read(), b'0.0.0.0 example.test\n')
|
|
result.output.close()
|
|
|
|
def test_no_supported_domain_returns_error(self):
|
|
with self.assertRaises(InvalidSourceError):
|
|
run_conversion(b'<!doctype html>\n<h1>Not found</h1>\n')
|
|
|
|
def test_invalid_encoding_returns_error(self):
|
|
with self.assertRaises(InvalidSourceError):
|
|
run_conversion(b'first.test\n\xff\xfe\n')
|
|
|
|
def test_binary_null_is_rejected(self):
|
|
with self.assertRaises(InvalidSourceError):
|
|
run_conversion(b'first.test\n\x00\x00\n')
|
|
|
|
def test_source_limit(self):
|
|
with self.assertRaises(SourceTooLargeError):
|
|
run_conversion(b'first.test\n', source_limit=5)
|
|
|
|
def test_response_content_length_limit(self):
|
|
with self.assertRaises(SourceTooLargeError):
|
|
run_conversion(b'first.test\n', headers={
|
|
"Content-Type": "application/octet-stream", "Content-Length": "2048"
|
|
}, source_limit=100)
|
|
|
|
def test_output_limit(self):
|
|
with self.assertRaises(SourceTooLargeError):
|
|
run_conversion(b'first.test\n', output_limit=8)
|
|
|
|
def test_overlong_line_is_skipped_without_buffering_it(self):
|
|
result = run_conversion(b'A' * 1000 + b'\nfirst.test\n', line_limit=32, chunk=4)
|
|
self.assertEqual(result.output.read(), b'0.0.0.0 first.test\n')
|
|
result.output.close()
|
|
|
|
def test_exception_from_upstream_happens_before_success_result(self):
|
|
with self.assertRaisesRegex(IOError, 'disconnected'):
|
|
run_conversion(b'first.test\nsecond.test\n', fail_after=7)
|
|
|
|
def test_debug_is_bytes(self):
|
|
result = run_conversion(b'first.test\n', debug=['# [DEBUG] start'])
|
|
output = result.output.read()
|
|
self.assertIn(b'# [DEBUG] start\n', output)
|
|
self.assertIn(b'0.0.0.0 first.test\n', output)
|
|
result.output.close()
|
|
|
|
def test_unrecognized_charset_rejected(self):
|
|
with self.assertRaises(InvalidSourceError):
|
|
run_conversion(b'first.test', headers={"Content-Type": "text/plain; charset=weird-charset-zzz"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|