OlaiEdge / tests /test_edge_engine.py

Upload folder using huggingface_hub

1aa277b verified 4 days ago

11.6 kB

"""
Comprehensive unit tests for FastAgent Edge Agent Engine.
Tests memory budgeting, multi-tab slot caching, GBNF grammar engine,
cryptographic offline licensing, and daemon HTTP endpoints.
"""
import json
import os
import time
import pytest

from fast_agent import (
    MemoryBudgetManager,
    MemoryPressureWarning,
    SlotCacheManager,
    SessionSlot,
    GrammarEngine,
    DESKTOP_SCHEMAS,
    LicenseManager,
    FastAgentLicense,
    LicenseTier,
    ModelManager,
    EDGE_MODEL_CATALOG,
    safe_parse_args,
)


# ============================================================================
# 1. Memory Budget and Crash Prevention Tests
# ============================================================================

def test_memory_budget_calculation():
    # Tight budget: 1GB budget with 800MB model -> forces q4_0, small context
    mgr_tight = MemoryBudgetManager(budget_mb=1024)
    params_tight = mgr_tight.compute_optimal_engine_params(model_size_mb=800, desired_context=2048)
    assert params_tight["ctk"] == "q4_0"
    assert params_tight["ctv"] == "q4_0"
    assert params_tight["context_size"] <= 1024

    # Ample budget: 4GB budget with 1GB model -> optimal q8_0, 2048 context
    mgr_ample = MemoryBudgetManager(budget_mb=4096)
    params_ample = mgr_ample.compute_optimal_engine_params(model_size_mb=1000, desired_context=4096)
    assert params_ample["ctk"] == "q8_0"
    assert params_ample["ctv"] == "q8_0"
    assert params_ample["context_size"] == 4096


def test_memory_snapshot_and_watchdog():
    mgr = MemoryBudgetManager(budget_mb=2048)
    snapshot = mgr.get_system_snapshot()
    assert "total_system_ram_mb" in snapshot
    assert "available_system_ram_mb" in snapshot
    assert snapshot["engine_budget_mb"] == 2048
    assert snapshot["total_system_ram_mb"] > 0


def test_circuit_breaker_on_memory_pressure():
    # Simulate an engine whose budget is heavily exceeded
    mgr = MemoryBudgetManager(budget_mb=10) # Unrealistic 10MB budget
    # The host process itself is > 10MB RSS
    with pytest.raises(MemoryPressureWarning) as exc_info:
        mgr.check_headroom_or_raise()
    assert "Memory budget exceeded" in str(exc_info.value) or "circuit breaker" in str(exc_info.value)


# ============================================================================
# 2. Multi-Tab Slot Caching Tests
# ============================================================================

def test_slot_cache_manager_tab_isolation():
    manager = SlotCacheManager(max_active_slots=3, default_history_turns=4)
    manager.set_pinned_system_prompt("System: You are an Edge Desktop Assistant.")

    # Tab 1: Code editor
    slot_code = manager.get_or_create_slot("tab-code-editor")
    slot_code.add_user_turn("Refactor function multiply()")
    slot_code.append_agent_turn("Thought: Refactoring.\nFinal Answer: Done.")

    # Tab 2: Note taking
    slot_notes = manager.get_or_create_slot("tab-notes")
    slot_notes.add_user_turn("Summarize chapter 1")
    slot_notes.append_agent_turn("Summary: Chapter 1 discusses edge AI.")

    # Verify tab isolation
    prompt_code = manager.build_slot_prompt("tab-code-editor")
    assert "Refactor function multiply()" in prompt_code
    assert "Summarize chapter 1" not in prompt_code
    assert "System: You are an Edge Desktop Assistant." in prompt_code

    prompt_notes = manager.build_slot_prompt("tab-notes")
    assert "Summarize chapter 1" in prompt_notes
    assert "Refactor function multiply()" not in prompt_notes


def test_slot_lru_eviction_and_restoration():
    manager = SlotCacheManager(max_active_slots=2)
    manager.get_or_create_slot("tab-1").add_user_turn("Hello 1")
    manager.get_or_create_slot("tab-2").add_user_turn("Hello 2")
    assert manager.active_slot_count() == 2

    # Accessing 3rd slot must evict LRU slot (tab-1) to archive
    manager.get_or_create_slot("tab-3").add_user_turn("Hello 3")
    assert manager.active_slot_count() == 2
    assert "tab-1" in manager.evicted_slots_archive

    # Re-accessing tab-1 must restore its conversation history from archive
    restored_tab1 = manager.get_or_create_slot("tab-1")
    assert any("Hello 1" in turn for turn in restored_tab1.conversation_history)


def test_slot_compaction_on_memory_pressure():
    manager = SlotCacheManager(max_active_slots=2, default_history_turns=6)
    slot = manager.get_or_create_slot("tab-long")
    for i in range(6):
        slot.add_user_turn(f"Turn {i}")
        slot.append_agent_turn(f"Response {i}")

    assert len(slot.conversation_history) == 6

    # Trigger emergency compaction
    manager.compact_all_slots(keep_last_n_turns=2)
    assert len(slot.conversation_history) == 2


# ============================================================================
# 3. Grammar Engine & Desktop Schemas Tests
# ============================================================================

def test_desktop_schemas_exist():
    for category in ["code_editor_patch", "note_summary", "cad_command", "accounting_entry"]:
        schema = GrammarEngine.get_desktop_schema(category)
        assert schema["type"] == "object"
        assert "properties" in schema
        assert "required" in schema


def test_json_schema_to_gbnf_conversion():
    schema = {
        "type": "object",
        "properties": {
            "action": {"type": "string", "enum": ["replace", "insert"]},
            "count": {"type": "integer"},
            "valid": {"type": "boolean"},
        },
        "required": ["action", "count"],
    }
    gbnf = GrammarEngine.json_schema_to_gbnf(schema)
    assert "root ::=" in gbnf
    assert '"\\"replace\\""' in gbnf
    assert '"\\"insert\\""' in gbnf
    assert "integer ::=" in gbnf


def test_grammar_validate_and_parse():
    schema = GrammarEngine.get_desktop_schema("code_editor_patch")

    valid_json = json.dumps({
        "file_path": "src/main.rs",
        "action": "replace",
        "start_line": 42,
        "replacement_code": "fn new_logic() {}",
        "commit_message": "Update logic"
    })
    ok, parsed, err = GrammarEngine.validate_and_parse(valid_json, schema)
    assert ok is True
    assert parsed["action"] == "replace"
    assert err is None

    # Missing required field 'replacement_code'
    invalid_json = json.dumps({
        "file_path": "src/main.rs",
        "action": "replace",
        "start_line": 42
    })
    ok, parsed, err = GrammarEngine.validate_and_parse(invalid_json, schema)
    assert ok is False
    assert "Missing required fields" in err


# ============================================================================
# 4. Licensing & Monetization Tests
# ============================================================================

def test_community_edition_license():
    lic = LicenseManager.verify_license(None)
    assert lic.tier == LicenseTier.COMMUNITY
    assert lic.is_valid is True
    assert lic.seat_count == 1
    assert lic.max_concurrent_slots >= 1


def test_commercial_license_token_generation_and_verification():
    token = LicenseManager.generate_token(
        licensee="Acme Desktop Software",
        tier=LicenseTier.COMMERCIAL,
        seat_count=50,
        max_slots=16,
        validity_days=30,
    )
    assert token.startswith("FA-COMM-")

    # Offline verification
    lic = LicenseManager.verify_license(token)
    assert lic.tier == LicenseTier.COMMERCIAL
    assert lic.licensee == "Acme Desktop Software"
    assert lic.seat_count == 50
    assert lic.max_concurrent_slots == 16
    assert lic.is_valid is True


def test_tampered_license_reverts_to_community():
    token = LicenseManager.generate_token(licensee="Legit Corp")
    tampered_token = token[:-4] + "dead"
    lic = LicenseManager.verify_license(tampered_token)
    assert lic.tier == LicenseTier.COMMUNITY
    assert lic.is_valid is False
    assert "mismatch" in lic.message.lower()


# ============================================================================
# 6. Daemon HTTP Server Route Tests (In-Memory Handler Verification)
# ============================================================================

def test_daemon_http_server_endpoints():
    import io
    from fast_agent_pkg.fast_agent.daemon import EdgeDaemonHandler

    class MockDaemon:
        def __init__(self):
            self.model_path = "gemma-2b-q4.gguf"
            self.slot_manager = SlotCacheManager(max_active_slots=4)
            self.license = LicenseManager.verify_license(None)

        def get_status(self):
            return {
                "status": "ready",
                "model_path": self.model_path,
                "memory": {"engine_rss_mb": 500, "engine_budget_mb": 2048},
                "slots": {"active_count": self.slot_manager.active_slot_count(), "max_allowed": 4},
                "license": self.license.to_dict(),
            }

    mock_daemon = MockDaemon()
    EdgeDaemonHandler.daemon_instance = mock_daemon

    def call_handler(method: str, path: str, body: dict = None):
        handler = EdgeDaemonHandler.__new__(EdgeDaemonHandler)
        handler.path = path
        handler.wfile = io.BytesIO()
        encoded_body = json.dumps(body).encode("utf-8") if body else b""
        handler.rfile = io.BytesIO(encoded_body)
        handler.headers = {"Content-Length": str(len(encoded_body))}

        status_codes = []
        sent_headers = {}

        def mock_send_response(code):
            status_codes.append(code)

        def mock_send_header(k, v):
            sent_headers[k] = v

        def mock_end_headers():
            pass

        handler.send_response = mock_send_response
        handler.send_header = mock_send_header
        handler.end_headers = mock_end_headers

        if method == "GET":
            handler.do_GET()
        elif method == "POST":
            handler.do_POST()

        raw_output = handler.wfile.getvalue().decode("utf-8")
        parsed = json.loads(raw_output) if raw_output else {}
        return status_codes[-1], parsed

    # 1. GET /health
    code, res = call_handler("GET", "/health")
    assert code == 200
    assert res["status"] == "ready"

    # 2. GET /v1/slots
    code, res = call_handler("GET", "/v1/slots")
    assert code == 200
    assert res["count"] == 0

    # 3. GET /v1/models
    code, res = call_handler("GET", "/v1/models")
    assert code == 200
    assert len(res["catalog"]) >= 4

    # 4. GET /v1/license
    code, res = call_handler("GET", "/v1/license")
    assert code == 200
    assert res["tier"] == "COMMUNITY"

    # 5. POST /v1/license/activate
    comm_token = LicenseManager.generate_token(licensee="Test Enterprise", seat_count=20)
    code, res = call_handler("POST", "/v1/license/activate", {"license_key": comm_token})
    assert code == 200
    assert res["license"]["tier"] == "COMMERCIAL"
    assert mock_daemon.license.tier == "COMMERCIAL"


def test_model_catalog_entries():
    catalog = ModelManager.list_catalog()
    assert len(catalog) >= 4
    model_ids = [m["id"] for m in catalog]
    assert "qwen2.5-1.5b" in model_ids
    assert "llama-3.2-1b" in model_ids
    assert "gemma-2b-q4" in model_ids


def test_gguf_magic_header_verification(tmp_path):
    # Valid GGUF file starts with b'GGUF'
    valid_file = tmp_path / "model.gguf"
    valid_file.write_bytes(b"GGUF\x03\x00\x00\x00" + b"\x00" * 100)
    assert ModelManager.verify_gguf_header(str(valid_file)) is True

    # Invalid header
    bad_file = tmp_path / "bad.gguf"
    bad_file.write_bytes(b"NOTGGUF" + b"\x00" * 100)
    assert ModelManager.verify_gguf_header(str(bad_file)) is False

	"""
	Comprehensive unit tests for FastAgent Edge Agent Engine.
	Tests memory budgeting, multi-tab slot caching, GBNF grammar engine,
	cryptographic offline licensing, and daemon HTTP endpoints.
	"""
	import json
	import os
	import time
	import pytest

	from fast_agent import (
	MemoryBudgetManager,
	MemoryPressureWarning,
	SlotCacheManager,
	SessionSlot,
	GrammarEngine,
	DESKTOP_SCHEMAS,
	LicenseManager,
	FastAgentLicense,
	LicenseTier,
	ModelManager,
	EDGE_MODEL_CATALOG,
	safe_parse_args,
	)


	# ============================================================================
	# 1. Memory Budget and Crash Prevention Tests
	# ============================================================================

	def test_memory_budget_calculation():
	# Tight budget: 1GB budget with 800MB model -> forces q4_0, small context
	mgr_tight = MemoryBudgetManager(budget_mb=1024)
	params_tight = mgr_tight.compute_optimal_engine_params(model_size_mb=800, desired_context=2048)
	assert params_tight["ctk"] == "q4_0"
	assert params_tight["ctv"] == "q4_0"
	assert params_tight["context_size"] <= 1024

	# Ample budget: 4GB budget with 1GB model -> optimal q8_0, 2048 context
	mgr_ample = MemoryBudgetManager(budget_mb=4096)
	params_ample = mgr_ample.compute_optimal_engine_params(model_size_mb=1000, desired_context=4096)
	assert params_ample["ctk"] == "q8_0"
	assert params_ample["ctv"] == "q8_0"
	assert params_ample["context_size"] == 4096


	def test_memory_snapshot_and_watchdog():
	mgr = MemoryBudgetManager(budget_mb=2048)
	snapshot = mgr.get_system_snapshot()
	assert "total_system_ram_mb" in snapshot
	assert "available_system_ram_mb" in snapshot
	assert snapshot["engine_budget_mb"] == 2048
	assert snapshot["total_system_ram_mb"] > 0


	def test_circuit_breaker_on_memory_pressure():
	# Simulate an engine whose budget is heavily exceeded
	mgr = MemoryBudgetManager(budget_mb=10) # Unrealistic 10MB budget
	# The host process itself is > 10MB RSS
	with pytest.raises(MemoryPressureWarning) as exc_info:
	mgr.check_headroom_or_raise()
	assert "Memory budget exceeded" in str(exc_info.value) or "circuit breaker" in str(exc_info.value)


	# ============================================================================
	# 2. Multi-Tab Slot Caching Tests
	# ============================================================================

	def test_slot_cache_manager_tab_isolation():
	manager = SlotCacheManager(max_active_slots=3, default_history_turns=4)
	manager.set_pinned_system_prompt("System: You are an Edge Desktop Assistant.")

	# Tab 1: Code editor
	slot_code = manager.get_or_create_slot("tab-code-editor")
	slot_code.add_user_turn("Refactor function multiply()")
	slot_code.append_agent_turn("Thought: Refactoring.\nFinal Answer: Done.")

	# Tab 2: Note taking
	slot_notes = manager.get_or_create_slot("tab-notes")
	slot_notes.add_user_turn("Summarize chapter 1")
	slot_notes.append_agent_turn("Summary: Chapter 1 discusses edge AI.")

	# Verify tab isolation
	prompt_code = manager.build_slot_prompt("tab-code-editor")
	assert "Refactor function multiply()" in prompt_code
	assert "Summarize chapter 1" not in prompt_code
	assert "System: You are an Edge Desktop Assistant." in prompt_code

	prompt_notes = manager.build_slot_prompt("tab-notes")
	assert "Summarize chapter 1" in prompt_notes
	assert "Refactor function multiply()" not in prompt_notes


	def test_slot_lru_eviction_and_restoration():
	manager = SlotCacheManager(max_active_slots=2)
	manager.get_or_create_slot("tab-1").add_user_turn("Hello 1")
	manager.get_or_create_slot("tab-2").add_user_turn("Hello 2")
	assert manager.active_slot_count() == 2

	# Accessing 3rd slot must evict LRU slot (tab-1) to archive
	manager.get_or_create_slot("tab-3").add_user_turn("Hello 3")
	assert manager.active_slot_count() == 2
	assert "tab-1" in manager.evicted_slots_archive

	# Re-accessing tab-1 must restore its conversation history from archive
	restored_tab1 = manager.get_or_create_slot("tab-1")
	assert any("Hello 1" in turn for turn in restored_tab1.conversation_history)


	def test_slot_compaction_on_memory_pressure():
	manager = SlotCacheManager(max_active_slots=2, default_history_turns=6)
	slot = manager.get_or_create_slot("tab-long")
	for i in range(6):
	slot.add_user_turn(f"Turn {i}")
	slot.append_agent_turn(f"Response {i}")

	assert len(slot.conversation_history) == 6

	# Trigger emergency compaction
	manager.compact_all_slots(keep_last_n_turns=2)
	assert len(slot.conversation_history) == 2


	# ============================================================================
	# 3. Grammar Engine & Desktop Schemas Tests
	# ============================================================================

	def test_desktop_schemas_exist():
	for category in ["code_editor_patch", "note_summary", "cad_command", "accounting_entry"]:
	schema = GrammarEngine.get_desktop_schema(category)
	assert schema["type"] == "object"
	assert "properties" in schema
	assert "required" in schema


	def test_json_schema_to_gbnf_conversion():
	schema = {
	"type": "object",
	"properties": {
	"action": {"type": "string", "enum": ["replace", "insert"]},
	"count": {"type": "integer"},
	"valid": {"type": "boolean"},
	},
	"required": ["action", "count"],
	}
	gbnf = GrammarEngine.json_schema_to_gbnf(schema)
	assert "root ::=" in gbnf
	assert '"\\"replace\\""' in gbnf
	assert '"\\"insert\\""' in gbnf
	assert "integer ::=" in gbnf


	def test_grammar_validate_and_parse():
	schema = GrammarEngine.get_desktop_schema("code_editor_patch")

	valid_json = json.dumps({
	"file_path": "src/main.rs",
	"action": "replace",
	"start_line": 42,
	"replacement_code": "fn new_logic() {}",
	"commit_message": "Update logic"
	})
	ok, parsed, err = GrammarEngine.validate_and_parse(valid_json, schema)
	assert ok is True
	assert parsed["action"] == "replace"
	assert err is None

	# Missing required field 'replacement_code'
	invalid_json = json.dumps({
	"file_path": "src/main.rs",
	"action": "replace",
	"start_line": 42
	})
	ok, parsed, err = GrammarEngine.validate_and_parse(invalid_json, schema)
	assert ok is False
	assert "Missing required fields" in err


	# ============================================================================
	# 4. Licensing & Monetization Tests
	# ============================================================================

	def test_community_edition_license():
	lic = LicenseManager.verify_license(None)
	assert lic.tier == LicenseTier.COMMUNITY
	assert lic.is_valid is True
	assert lic.seat_count == 1
	assert lic.max_concurrent_slots >= 1


	def test_commercial_license_token_generation_and_verification():
	token = LicenseManager.generate_token(
	licensee="Acme Desktop Software",
	tier=LicenseTier.COMMERCIAL,
	seat_count=50,
	max_slots=16,
	validity_days=30,
	)
	assert token.startswith("FA-COMM-")

	# Offline verification
	lic = LicenseManager.verify_license(token)
	assert lic.tier == LicenseTier.COMMERCIAL
	assert lic.licensee == "Acme Desktop Software"
	assert lic.seat_count == 50
	assert lic.max_concurrent_slots == 16
	assert lic.is_valid is True


	def test_tampered_license_reverts_to_community():
	token = LicenseManager.generate_token(licensee="Legit Corp")
	tampered_token = token[:-4] + "dead"
	lic = LicenseManager.verify_license(tampered_token)
	assert lic.tier == LicenseTier.COMMUNITY
	assert lic.is_valid is False
	assert "mismatch" in lic.message.lower()


	# ============================================================================
	# 6. Daemon HTTP Server Route Tests (In-Memory Handler Verification)
	# ============================================================================

	def test_daemon_http_server_endpoints():
	import io
	from fast_agent_pkg.fast_agent.daemon import EdgeDaemonHandler

	class MockDaemon:
	def __init__(self):
	self.model_path = "gemma-2b-q4.gguf"
	self.slot_manager = SlotCacheManager(max_active_slots=4)
	self.license = LicenseManager.verify_license(None)

	def get_status(self):
	return {
	"status": "ready",
	"model_path": self.model_path,
	"memory": {"engine_rss_mb": 500, "engine_budget_mb": 2048},
	"slots": {"active_count": self.slot_manager.active_slot_count(), "max_allowed": 4},
	"license": self.license.to_dict(),
	}

	mock_daemon = MockDaemon()
	EdgeDaemonHandler.daemon_instance = mock_daemon

	def call_handler(method: str, path: str, body: dict = None):
	handler = EdgeDaemonHandler.__new__(EdgeDaemonHandler)
	handler.path = path
	handler.wfile = io.BytesIO()
	encoded_body = json.dumps(body).encode("utf-8") if body else b""
	handler.rfile = io.BytesIO(encoded_body)
	handler.headers = {"Content-Length": str(len(encoded_body))}

	status_codes = []
	sent_headers = {}

	def mock_send_response(code):
	status_codes.append(code)

	def mock_send_header(k, v):
	sent_headers[k] = v

	def mock_end_headers():
	pass

	handler.send_response = mock_send_response
	handler.send_header = mock_send_header
	handler.end_headers = mock_end_headers

	if method == "GET":
	handler.do_GET()
	elif method == "POST":
	handler.do_POST()

	raw_output = handler.wfile.getvalue().decode("utf-8")
	parsed = json.loads(raw_output) if raw_output else {}
	return status_codes[-1], parsed

	# 1. GET /health
	code, res = call_handler("GET", "/health")
	assert code == 200
	assert res["status"] == "ready"

	# 2. GET /v1/slots
	code, res = call_handler("GET", "/v1/slots")
	assert code == 200
	assert res["count"] == 0

	# 3. GET /v1/models
	code, res = call_handler("GET", "/v1/models")
	assert code == 200
	assert len(res["catalog"]) >= 4

	# 4. GET /v1/license
	code, res = call_handler("GET", "/v1/license")
	assert code == 200
	assert res["tier"] == "COMMUNITY"

	# 5. POST /v1/license/activate
	comm_token = LicenseManager.generate_token(licensee="Test Enterprise", seat_count=20)
	code, res = call_handler("POST", "/v1/license/activate", {"license_key": comm_token})
	assert code == 200
	assert res["license"]["tier"] == "COMMERCIAL"
	assert mock_daemon.license.tier == "COMMERCIAL"


	def test_model_catalog_entries():
	catalog = ModelManager.list_catalog()
	assert len(catalog) >= 4
	model_ids = [m["id"] for m in catalog]
	assert "qwen2.5-1.5b" in model_ids
	assert "llama-3.2-1b" in model_ids
	assert "gemma-2b-q4" in model_ids


	def test_gguf_magic_header_verification(tmp_path):
	# Valid GGUF file starts with b'GGUF'
	valid_file = tmp_path / "model.gguf"
	valid_file.write_bytes(b"GGUF\x03\x00\x00\x00" + b"\x00" * 100)
	assert ModelManager.verify_gguf_header(str(valid_file)) is True

	# Invalid header
	bad_file = tmp_path / "bad.gguf"
	bad_file.write_bytes(b"NOTGGUF" + b"\x00" * 100)
	assert ModelManager.verify_gguf_header(str(bad_file)) is False