test_suite.py

#!/usr/bin/env python3
"""Automated test suite for Viper transpiler.

Usage: python3 test_suite.py
Runs all tests and reports pass/fail for each feature.
"""

import ast
import subprocess
import sys
import tempfile
import os
from typing import Dict, List, Tuple

from core import Transpiler, TranspileError
from resolver import ImportResolver


def transpile_and_run(py_src: str) -> Tuple[bool, str, str]:
    """Transpile Python source, compile to C++, run, and return (ok, stdout, stderr)."""
    try:
        tree = ast.parse(py_src)
    except SyntaxError as e:
        return (False, "", f"Python syntax error: {e}")

    r = ImportResolver({"test"})
    try:
        extra = r.collect([tree])
    except TranspileError as e:
        return (False, "", f"Import resolution error: {e}")

    modules = [tree]
    names = ["test"]
    for t, n in extra:
        modules.append(t)
        names.append(n)

    t = Transpiler(modules, names)
    try:
        cpp = t.transpile()
    except TranspileError as e:
        return (False, "", f"Transpile error: {e}")

    # Check for skipped functions
    if "// [Viper skip]" in cpp:
        skips = [l.strip() for l in cpp.split("\n") if "// [Viper skip]" in l]
        return (False, "", "\n".join(skips))

    # Compile and run
    with tempfile.NamedTemporaryFile(suffix=".cpp", delete=False, mode="w") as f:
        f.write(cpp)
        cpp_path = f.name

    exe_path = cpp_path[:-4]
    try:
        compile_result = subprocess.run(
            ["g++", "-std=c++17", "-O2", "-o", exe_path, cpp_path],
            capture_output=True, text=True, timeout=30
        )
        if compile_result.returncode != 0:
            return (False, "", f"Compilation error:\n{compile_result.stderr[:500]}")

        run_result = subprocess.run(
            [exe_path], capture_output=True, text=True, timeout=10
        )
        return (True, run_result.stdout.strip(), run_result.stderr.strip())
    finally:
        try:
            os.unlink(cpp_path)
            os.unlink(exe_path)
        except OSError:
            pass


# ============================================================
# Test cases: (name, python_source_with_main, expected_output)
# ============================================================

TESTS: List[Tuple[str, str, str]] = [
    # --- Primitives ---
    ("int arithmetic", """
def f() -> int:
    return 2 + 3 * 4
if __name__ == "__main__": print(f())
""", "14"),

    ("float arithmetic", """
def f() -> float:
    return 10.0 / 4.0
if __name__ == "__main__": print(f())
""", "2.5"),

    ("bool operations", """
def f() -> bool:
    return True and not False
if __name__ == "__main__": print(f())
""", "true"),

    ("string concat", """
def f() -> str:
    a: str = "Hello, "
    b: str = "World!"
    return a + b
if __name__ == "__main__": print(f())
""", "Hello, World!"),

    ("string upper", """
def f() -> str:
    return "hello".upper()
if __name__ == "__main__": print(f())
""", "HELLO"),

    # --- Collections ---
    ("list indexing", """
def f() -> int:
    lst: list[int] = [10, 20, 30]
    return lst[0] + lst[2]
if __name__ == "__main__": print(f())
""", "40"),

    ("list append", """
def f() -> int:
    lst: list[int] = [1, 2]
    lst.append(3)
    return len(lst)
if __name__ == "__main__": print(f())
""", "3"),

    ("list slice", """
def f() -> list[int]:
    lst: list[int] = [0, 1, 2, 3, 4]
    return lst[1:4]
if __name__ == "__main__": print(f())
""", "[1, 2, 3]"),

    ("dict operations", """
def f() -> int:
    d: dict[str, int] = {}
    d["x"] = 10
    d["y"] = 20
    return d["x"] + d["y"]
if __name__ == "__main__": print(f())
""", "30"),

    ("set union", """
def f() -> int:
    a: set[int] = {1, 2, 3}
    b: set[int] = {3, 4, 5}
    return len(a | b)
if __name__ == "__main__": print(f())
""", "5"),

    ("tuple", """
def f() -> float:
    t: tuple[int, float, str] = (1, 3.14, "pi")
    return t[1]
if __name__ == "__main__": print(f())
""", "3.14"),

    # --- Control flow ---
    ("if else", """
def f(x: int) -> str:
    if x > 0: return "pos"
    elif x < 0: return "neg"
    else: return "zero"
if __name__ == "__main__": print(f(5))
""", "pos"),

    ("while loop", """
def f() -> int:
    total: int = 0
    i: int = 0
    while i < 5:
        total = total + i
        i = i + 1
    return total
if __name__ == "__main__": print(f())
""", "10"),

    ("for range", """
def f() -> int:
    total: int = 0
    for i in range(1, 6):
        total = total + i
    return total
if __name__ == "__main__": print(f())
""", "15"),

    ("for list", """
def f() -> int:
    total: int = 0
    nums: list[int] = [10, 20, 30]
    for n in nums:
        total = total + n
    return total
if __name__ == "__main__": print(f())
""", "60"),

    ("ternary", """
def f(x: int) -> str:
    return "yes" if x > 0 else "no"
if __name__ == "__main__": print(f(1))
""", "yes"),

    # --- Functions ---
    ("default param", """
def greet(name: str = "World") -> str:
    return "Hello, " + name
if __name__ == "__main__": print(greet())
""", "Hello, World"),

    ("keyword args", """
def order(x: int, y: int = 0) -> int:
    return x * 10 + y
if __name__ == "__main__": print(order(y=5, x=3))
""", "35"),

    ("varargs", """
def sum_all(*args: int) -> int:
    total: int = 0
    for n in args:
        total = total + n
    return total
if __name__ == "__main__": print(sum_all(1, 2, 3, 4))
""", "10"),

    ("generator yield", """
def seq(n: int) -> list[int]:
    i: int = 0
    while i < n:
        yield i
        i = i + 1
if __name__ == "__main__": print(len(seq(5)))
""", "5"),

    # --- Classes ---
    ("class basic", """
class Point:
    x: float
    y: float
    def __init__(self, x: float, y: float) -> None:
        self.x = x
        self.y = y
    def sum(self) -> float:
        return self.x + self.y
def f() -> float:
    p: Point = Point(3.0, 4.0)
    return p.sum()
if __name__ == "__main__": print(f())
""", "7"),

    ("inheritance", """
class Animal:
    name: str
    def __init__(self, name: str) -> None:
        self.name = name
    def speak(self) -> str:
        return self.name
class Dog(Animal):
    def speak(self) -> str:
        return self.name + " barks"
def f() -> str:
    return Dog("Rex").speak()
if __name__ == "__main__": print(f())
""", "Rex barks"),

    ("dataclass", """
from dataclasses import dataclass
@dataclass
class Point:
    x: float
    y: float
def f() -> str:
    p: Point = Point(1.0, 2.0)
    return str(p)
if __name__ == "__main__": print(f())
""", "Point(x=1.000000, y=2.000000)"),

    ("enum", """
from enum import IntEnum
class Color(IntEnum):
    RED = 1
    GREEN = 2
def f() -> int:
    return Color.RED + Color.GREEN
if __name__ == "__main__": print(f())
""", "3"),

    ("operator overload", """
class Vec:
    x: float
    def __init__(self, x: float) -> None:
        self.x = x
    def __add__(self, other: Vec) -> Vec:
        return Vec(self.x + other.x)
    def __str__(self) -> str:
        return str(self.x)
def f() -> float:
    a: Vec = Vec(3.0)
    b: Vec = Vec(4.0)
    return (a + b).x
if __name__ == "__main__": print(f())
""", "7"),

    # --- Exception handling ---
    ("try except", """
def f() -> str:
    try:
        raise ValueError("test")
    except ValueError:
        return "caught"
    return "not caught"
if __name__ == "__main__": print(f())
""", "caught"),

    # --- Lambda ---
    ("lambda typed", """
from typing import Callable
def f() -> int:
    fn: Callable[[int], int] = lambda x: x * 3
    return fn(7)
if __name__ == "__main__": print(f())
""", "21"),

    # --- Comprehensions ---
    ("list comprehension", """
def f() -> list[int]:
    return [x * x for x in range(1, 4)]
if __name__ == "__main__": print(f())
""", "[1, 4, 9]"),

    # --- Built-ins ---
    ("sum", """
def f() -> int:
    return sum([1, 2, 3, 4])
if __name__ == "__main__": print(f())
""", "10"),

    ("sorted", """
def f() -> list[int]:
    return sorted([3, 1, 2])
if __name__ == "__main__": print(f())
""", "[1, 2, 3]"),

    ("map", """
def square(x: int) -> int:
    return x * x
def f() -> list[int]:
    return list(map(square, [1, 2, 3]))
if __name__ == "__main__": print(f())
""", "[1, 4, 9]"),

    ("filter", """
def pos(x: int) -> bool:
    return x > 0
def f() -> list[int]:
    return list(filter(pos, [-1, 0, 1, 2]))
if __name__ == "__main__": print(f())
""", "[1, 2]"),

    ("any all", """
def f() -> bool:
    return any([False, True]) and all([True, True])
if __name__ == "__main__": print(f())
""", "true"),

    # --- Bytes ---
    ("bytes literal", """
def f() -> int:
    b: bytes = b"ABC"
    return b[0]
if __name__ == "__main__": print(f())
""", "65"),

    ("bytes concat", """
def f() -> int:
    a: bytes = b"AB"
    b: bytes = b"CD"
    c: bytes = a + b
    return len(c)
if __name__ == "__main__": print(f())
""", "4"),

    # --- New features ---
    ("kwargs", """
def log(msg: str, **opts) -> str:
    return msg
def f() -> str:
    return log("test")
if __name__ == "__main__": print(f())
""", "test"),

    ("strenum", """
from enum import Enum
class Color(Enum):
    RED = "red"
    GREEN = "green"
def f() -> str:
    return Color.RED
if __name__ == "__main__": print(f())
""", "red"),

    ("slice assign", """
def f() -> list[int]:
    lst: list[int] = [1, 2, 3, 4]
    lst[1:3] = [10, 20]
    return lst
if __name__ == "__main__": print(f())
""", "[1, 10, 20, 4]"),

    ("ordered dict", """
from collections import OrderedDict
def f() -> int:
    d: OrderedDict[str, int] = OrderedDict()
    d["x"] = 10
    return d["x"]
if __name__ == "__main__": print(f())
""", "10"),

    ("counter", """
from collections import Counter
def f() -> dict[str, int]:
    return Counter(["a", "b", "a"])
if __name__ == "__main__": print(f())
""", "{b: 1, a: 2}"),

    ("any type", """
from typing import Any, cast
def f() -> int:
    a: Any = 42
    return cast(int, a)
if __name__ == "__main__": print(f())
""", "42"),

    ("dataclass default", """
from dataclasses import dataclass
@dataclass
class Cfg:
    x: float = 3.5
def f() -> float:
    c: Cfg = Cfg()
    return c.x
if __name__ == "__main__": print(f())
""", "3.5"),

    ("dataclass frozen", """
from dataclasses import dataclass
@dataclass(frozen=True)
class Pt:
    x: float
    y: float
def f() -> float:
    p: Pt = Pt(3.0, 4.0)
    return p.x + p.y
if __name__ == "__main__": print(f())
""", "7"),

    ("polymorphism", """
class Animal:
    name: str
    def __init__(self, name: str) -> None:
        self.name = name
    def speak(self) -> str:
        return "generic"
class Dog(Animal):
    def speak(self) -> str:
        return self.name + " barks"
def process(a: Animal) -> str:
    return a.speak()
def f() -> str:
    return process(Dog("Fido"))
if __name__ == "__main__": print(f())
""", "Fido barks"),

    ("walrus", """
def f() -> int:
    result: int = 0
    i: int = 0
    while (i := i + 1) <= 3:
        result = result + i
    return result
if __name__ == "__main__": print(f())
""", "6"),

    ("for else", """
def f() -> str:
    for i in range(3):
        if i == 99:
            break
    else:
        return "not found"
    return "??"
if __name__ == "__main__": print(f())
""", "not found"),

    ("list repeat", """
def f() -> list[int]:
    return [1, 2] * 3
if __name__ == "__main__": print(f())
""", "[1, 2, 1, 2, 1, 2]"),

    ("list reverse slice", """
def f() -> list[int]:
    lst: list[int] = [1, 2, 3, 4, 5]
    return lst[::-1]
if __name__ == "__main__": print(f())
""", "[5, 4, 3, 2, 1]"),

    ("negative index", """
def f() -> int:
    nums: list[int] = [10, 20, 30, 40]
    return nums[-1] + nums[-2]
if __name__ == "__main__": print(f())
""", "70"),

    ("match or pattern", """
def f(x: int) -> str:
    match x:
        case 1 | 2:
            return "one or two"
        case _:
            return "other"
if __name__ == "__main__": print(f(2))
""", "one or two"),

    ("string split", """
def f() -> int:
    s: str = "a,b,c"
    parts: list[str] = s.split(",")
    return len(parts)
if __name__ == "__main__": print(f())
""", "3"),
]


def main():
    passed = 0
    failed = 0
    for name, py_src, expected in TESTS:
        ok, stdout, stderr = transpile_and_run(py_src)
        if ok and stdout == expected:
            passed += 1
            print(f"  PASS  {name}")
        else:
            failed += 1
            print(f"  FAIL  {name}")
            if stderr:
                for line in stderr.split("\n")[:3]:
                    print(f"        {line}")
            if stdout and stdout != expected:
                print(f"        expected: {expected}")
                print(f"        got:      {stdout}")

    print()
    print(f"Results: {passed} passed, {failed} failed, {passed + failed} total")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    sys.exit(main())