fix

tests/test_source_conversion.py

@@ -0,0 +1,122 @@
+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()