mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
131 lines
3.6 KiB
Python
131 lines
3.6 KiB
Python
"""Shared pytest fixtures for the Obliteratus test suite."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import socket
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
|
|
_EXTERNAL_MARKERS = ("network", "download", "remote")
|
|
_EVIDENCE_MARKERS = (
|
|
"cpu",
|
|
"integration",
|
|
"slow",
|
|
"gpu",
|
|
"mps",
|
|
"mlx",
|
|
"network",
|
|
"download",
|
|
"remote",
|
|
"operator_ui",
|
|
)
|
|
|
|
|
|
def pytest_collection_modifyitems(items):
|
|
"""Attach stable test-layer markers to JUnit timing evidence."""
|
|
for item in items:
|
|
markers = sorted(
|
|
name for name in _EVIDENCE_MARKERS
|
|
if item.get_closest_marker(name) is not None
|
|
)
|
|
item.user_properties.append(("duration_markers", ",".join(markers) or "unmarked"))
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def offline_test_environment(request, monkeypatch):
|
|
"""Fail accidental network access and force offline Hugging Face behavior."""
|
|
if any(request.node.get_closest_marker(name) for name in _EXTERNAL_MARKERS):
|
|
return
|
|
|
|
monkeypatch.setenv("HF_DATASETS_OFFLINE", "1")
|
|
monkeypatch.setenv("HF_HUB_DISABLE_TELEMETRY", "1")
|
|
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
|
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
|
|
|
def reject_network(*_args, **_kwargs):
|
|
raise RuntimeError(
|
|
"unmarked tests may not access the network; add an explicit "
|
|
"network, download, or remote marker",
|
|
)
|
|
|
|
monkeypatch.setattr(socket, "create_connection", reject_network)
|
|
monkeypatch.setattr(socket.socket, "connect", reject_network)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_model():
|
|
"""A minimal mock transformer model.
|
|
|
|
Provides:
|
|
- model.config with config.num_hidden_layers = 4
|
|
- model.named_parameters() returning fake weight tensors
|
|
"""
|
|
import torch
|
|
|
|
model = MagicMock()
|
|
|
|
# Config with num_hidden_layers
|
|
config = MagicMock()
|
|
config.num_hidden_layers = 4
|
|
model.config = config
|
|
|
|
# named_parameters returns fake weight tensors across 4 layers
|
|
fake_params = []
|
|
for layer_idx in range(4):
|
|
weight = torch.randn(768, 768)
|
|
fake_params.append((f"model.layers.{layer_idx}.self_attn.q_proj.weight", weight))
|
|
fake_params.append((f"model.layers.{layer_idx}.self_attn.v_proj.weight", weight))
|
|
fake_params.append((f"model.layers.{layer_idx}.mlp.gate_proj.weight", weight))
|
|
model.named_parameters.return_value = fake_params
|
|
|
|
return model
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_tokenizer():
|
|
"""A minimal mock tokenizer with encode, decode, and apply_chat_template."""
|
|
tokenizer = MagicMock()
|
|
|
|
tokenizer.encode.return_value = [1, 2, 3, 4, 5]
|
|
tokenizer.decode.return_value = "Hello, this is a decoded string."
|
|
tokenizer.apply_chat_template.return_value = [1, 2, 3, 4, 5, 6, 7]
|
|
|
|
tokenizer.pad_token = "<pad>"
|
|
tokenizer.eos_token = "<eos>"
|
|
|
|
return tokenizer
|
|
|
|
|
|
@pytest.fixture
|
|
def refusal_direction():
|
|
"""A normalized random torch tensor of shape (768,)."""
|
|
import torch
|
|
|
|
t = torch.randn(768)
|
|
return t / t.norm()
|
|
|
|
|
|
@pytest.fixture
|
|
def activation_pair():
|
|
"""A tuple of (harmful_activations, harmless_activations) as random tensors of shape (10, 768)."""
|
|
import torch
|
|
|
|
harmful_activations = torch.randn(10, 768)
|
|
harmless_activations = torch.randn(10, 768)
|
|
return (harmful_activations, harmless_activations)
|
|
|
|
|
|
@pytest.fixture
|
|
def tmp_output_dir(tmp_path):
|
|
"""A clean temporary output directory for test artifacts."""
|
|
output_dir = tmp_path / "test_output"
|
|
output_dir.mkdir()
|
|
return output_dir
|