"""Codex compresses request bodies with zstd; the bridge must be able to read them."""
import ast
import unittest
from pathlib import Path

tree = ast.parse(Path(__file__).with_name("bridge.py").read_text(encoding="utf-8"))
wanted = {"_zstd_decompress", "_zstd_ready", "decompress"}
env = {}
exec(compile(ast.Module(body=[n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in wanted],
                        type_ignores=[]), "bridge.py", "exec"), env)


class Zstd(unittest.TestCase):
    def test_self_check_matches_reality(self):
        try:
            try:
                from compression import zstd as z
                packed = z.compress(b"govde" * 50)
            except ImportError:
                import zstandard
                packed = zstandard.ZstdCompressor().compress(b"govde" * 50)
        except ImportError:
            self.assertFalse(env["_zstd_ready"]())      # no zstd here: the bridge must say so
            return
        self.assertTrue(env["_zstd_ready"]())
        self.assertEqual(env["decompress"](packed, "zstd"), b"govde" * 50)

    def test_other_encodings(self):
        import gzip
        self.assertEqual(env["decompress"](gzip.compress(b"x"), "gzip"), b"x")
        self.assertEqual(env["decompress"](b"duz", None), b"duz")


if __name__ == "__main__":
    unittest.main()