conftest.py

"""
Shared pytest fixtures for the DNN decompiler test suite.
"""
import pytest
import tempfile
import os
import shutil
from pathlib import Path
from unittest.mock import Mock, MagicMock
import numpy as np


@pytest.fixture
def temp_dir():
    """Create a temporary directory for tests."""
    temp_path = tempfile.mkdtemp()
    yield Path(temp_path)
    shutil.rmtree(temp_path, ignore_errors=True)


@pytest.fixture
def sample_binary_path(temp_dir):
    """Create a sample binary file for testing."""
    binary_path = temp_dir / "test_binary.bin"
    with open(binary_path, "wb") as f:
        f.write(b"\x00" * 1024)  # Simple 1KB binary
    return binary_path


@pytest.fixture
def mock_angr_project():
    """Mock angr project for testing."""
    project = Mock()
    project.arch = Mock()
    project.arch.name = "ARM"
    project.loader = Mock()
    project.loader.main_object = Mock()
    project.funcs = {}
    project.outer_loops = {}
    project.func_calling_regs = {}
    return project


@pytest.fixture
def mock_capstone():
    """Mock capstone disassembler."""
    mock_cs = Mock()
    mock_cs.disasm = Mock(return_value=[])
    return mock_cs


@pytest.fixture
def sample_numpy_array():
    """Sample numpy array for testing."""
    return np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)


@pytest.fixture
def mock_onnx_model():
    """Mock ONNX model for testing."""
    model = Mock()
    model.graph = Mock()
    model.graph.node = []
    model.graph.input = []
    model.graph.output = []
    return model


@pytest.fixture
def test_config():
    """Test configuration dictionary."""
    return {
        "timeout": 10,
        "debug": True,
        "arch": "ARM",
        "endian": "little"
    }


@pytest.fixture
def mock_memory_record():
    """Mock memory record for testing."""
    record = Mock()
    record.addr = 0x1000
    record.size = 4
    record.data = b"\x01\x02\x03\x04"
    return record


@pytest.fixture
def sample_lifted_ast():
    """Sample lifted AST for testing."""
    from src.lifted_ast import LiftedAST, AST_OP
    ast = LiftedAST(None, None, None, None)
    ast.op_type = AST_OP.CONV
    return ast


@pytest.fixture(autouse=True)
def cleanup_temp_files():
    """Auto-cleanup fixture to remove temporary files after each test."""
    yield
    # Clean up any .tmp files or test artifacts
    for tmp_file in Path("/tmp").glob("test_*"):
        try:
            if tmp_file.is_file():
                tmp_file.unlink()
            elif tmp_file.is_dir():
                shutil.rmtree(tmp_file)
        except (OSError, PermissionError):
            pass


@pytest.fixture
def capture_logs(caplog):
    """Fixture to capture and provide access to log messages."""
    import logging
    caplog.set_level(logging.DEBUG)
    return caplog


# Custom markers
def pytest_configure(config):
    """Register custom markers."""
    config.addinivalue_line("markers", "unit: mark test as a unit test")
    config.addinivalue_line("markers", "integration: mark test as an integration test")
    config.addinivalue_line("markers", "slow: mark test as slow running")