"""
tests/test_token_budget_slots.py — Tests for slot-based allocation and bundle compression.

Covers Issues #374 (slot allocation) and #373 (bundle compression):
  - allocate_slots(): priority-ordered slot filling, overflow trimming, edge cases
  - compress_to_fit(): compression level selection, budget pressure, edge cases
  - context_window_budgets(): fraction-based context-window splitting
  - Config fields: memory_fraction, user_fraction, reserve_fraction
  - Backward compat: existing allocate_layers() unchanged
"""

from __future__ import annotations

import pytest

from cortex.core.token_budget import TokenBudgetManager
from cortex.types import (
    BudgetAllocation,
    CompressedItem,
    CompressionLevel,
    InjectionSlot,
    MemoryLayer,
)


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def manager() -> TokenBudgetManager:
    """Default manager: 4000 total, 800 cornerstone."""
    return TokenBudgetManager()


@pytest.fixture
def small_manager() -> TokenBudgetManager:
    """Small manager for overflow/truncation tests."""
    return TokenBudgetManager(total_budget=200, cornerstone_budget=50)


# Token-stable test strings.  Token counts verified against tiktoken cl100k_base.
# Each string here is 2 tokens: "hello world" → [15, 1917]
ITEM_2T = "hello world"
# 'old memory' → 2 tokens
ITEM_OLD = "old memory"
# 21-token string: 'word ' * 20 → 21 tokens (20 × 'word' + space run)
ITEM_21T = "word " * 20


# ---------------------------------------------------------------------------
# InjectionSlot enum helpers (sanity)
# ---------------------------------------------------------------------------

class TestInjectionSlotEnum:
    def test_values(self):
        assert InjectionSlot.TASK_CONTEXT.value == "task_context"
        assert InjectionSlot.EPISODIC.value == "episodic"

    def test_priority_ordering(self):
        ordered = InjectionSlot.by_priority()
        assert ordered[0] == InjectionSlot.TASK_CONTEXT
        assert ordered[-1] == InjectionSlot.EPISODIC

    def test_priority_values_strictly_increasing(self):
        ordered = InjectionSlot.by_priority()
        priorities = [s.priority for s in ordered]
        assert priorities == sorted(priorities)

    def test_default_priority_names(self):
        names = InjectionSlot.default_priority_names()
        assert names[0] == "task_context"
        assert names[-1] == "episodic"
        assert len(names) == 6


# ---------------------------------------------------------------------------
# allocate_slots() — basic behaviour
# ---------------------------------------------------------------------------

class TestAllocateSlots:
    # --- Edge cases ---

    def test_zero_budget_returns_empty(self, manager):
        slots = {InjectionSlot.TASK_CONTEXT: [ITEM_2T]}
        result = manager.allocate_slots(0, slots)
        assert result == {}

    def test_negative_budget_returns_empty(self, manager):
        slots = {InjectionSlot.TASK_CONTEXT: [ITEM_2T]}
        result = manager.allocate_slots(-100, slots)
        assert result == {}

    def test_empty_slots_dict_returns_empty(self, manager):
        result = manager.allocate_slots(1000, {})
        assert result == {}

    def test_slot_with_no_items_is_omitted(self, manager):
        slots = {
            InjectionSlot.TASK_CONTEXT: [],
            InjectionSlot.EPISODIC: [ITEM_2T],
        }
        result = manager.allocate_slots(100, slots)
        assert "task_context" not in result
        assert "episodic" in result

    # --- Single slot ---

    def test_single_slot_all_items_fit(self, manager):
        items = [ITEM_2T, ITEM_2T]  # 4 tokens total
        result = manager.allocate_slots(10, {InjectionSlot.TASK_CONTEXT: items})
        assert result == {"task_context": 4}

    def test_single_slot_partial_items_fit(self, manager):
        # First item: 2 tokens; second would push total to 4 but budget=3
        items = [ITEM_2T, ITEM_2T]
        result = manager.allocate_slots(3, {InjectionSlot.TASK_CONTEXT: items})
        # Greedy: first item fits (2), second doesn't (2+2=4 > 3) → break
        assert result == {"task_context": 2}

    def test_single_slot_nothing_fits(self, manager):
        # Each item is 2 tokens, budget is only 1
        result = manager.allocate_slots(1, {InjectionSlot.TASK_CONTEXT: [ITEM_2T]})
        assert result == {}

    # --- Return type and keys ---

    def test_returns_string_keys(self, manager):
        result = manager.allocate_slots(100, {InjectionSlot.TASK_CONTEXT: [ITEM_2T]})
        for key in result:
            assert isinstance(key, str), f"Expected str key, got {type(key)}"

    def test_values_are_int(self, manager):
        result = manager.allocate_slots(100, {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T],
            InjectionSlot.EPISODIC: [ITEM_OLD],
        })
        for v in result.values():
            assert isinstance(v, int)

    # --- Priority ordering ---

    def test_high_priority_wins_over_low_priority(self, manager):
        """TASK_CONTEXT (priority 0) must be filled before EPISODIC (priority 5)."""
        slots = {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T],   # 2 tokens
            InjectionSlot.EPISODIC: [ITEM_OLD],       # 2 tokens
        }
        # Budget=3: only the first slot (2 tokens) fits; second needs 2 more but 1 remains
        result = manager.allocate_slots(3, slots)
        assert "task_context" in result
        assert "episodic" not in result

    def test_both_slots_fit_within_budget(self, manager):
        slots = {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T],   # 2 tokens
            InjectionSlot.EPISODIC: [ITEM_OLD],       # 2 tokens
        }
        result = manager.allocate_slots(4, slots)
        assert result["task_context"] == 2
        assert result["episodic"] == 2

    def test_priority_ordering_all_six_slots(self, manager):
        """With budget that fits only two slots, highest two priorities win."""
        slots = {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T],
            InjectionSlot.BOOT_CONTEXT: [ITEM_2T],
            InjectionSlot.OPEN_LOOPS: [ITEM_2T],
            InjectionSlot.SEMANTIC_CLAIMS: [ITEM_2T],
            InjectionSlot.PROCEDURAL: [ITEM_2T],
            InjectionSlot.EPISODIC: [ITEM_2T],
        }
        result = manager.allocate_slots(4, slots)
        assert "task_context" in result
        assert "boot_context" in result
        # Lower-priority slots should NOT appear (only 4 token budget → 2 slots × 2 tokens)
        for low in ("open_loops", "semantic_claims", "procedural", "episodic"):
            assert low not in result, f"Expected {low!r} to be starved"

    def test_total_allocated_does_not_exceed_budget(self, manager):
        slots = {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T] * 5,
            InjectionSlot.BOOT_CONTEXT: [ITEM_2T] * 5,
            InjectionSlot.EPISODIC: [ITEM_OLD] * 5,
        }
        budget = 13
        result = manager.allocate_slots(budget, slots)
        total = sum(result.values())
        assert total <= budget

    # --- Multiple items per slot ---

    def test_multiple_items_per_slot_summed(self, manager):
        """Three 2-token items in one slot → 6 tokens allocated."""
        items = [ITEM_2T, ITEM_2T, ITEM_2T]
        result = manager.allocate_slots(100, {InjectionSlot.TASK_CONTEXT: items})
        assert result["task_context"] == 6

    def test_greedy_stops_at_first_non_fitting_item(self, manager):
        """If item N doesn't fit, items N+1, N+2, ... are also skipped."""
        # item 1: 2 tokens, item 2: 21 tokens (too big), item 3: 2 tokens (would fit but skipped)
        items = [ITEM_2T, ITEM_21T, ITEM_2T]
        result = manager.allocate_slots(10, {InjectionSlot.TASK_CONTEXT: items})
        # item1 fits (2), item2 doesn't (2+21=23 > 10) → breaks → item3 never tried
        assert result["task_context"] == 2

    # --- Custom slot priorities ---

    def test_custom_slot_priorities_reorders_filling(self):
        """EPISODIC-first priority order means episodic wins tight budget."""
        m = TokenBudgetManager(
            slot_priorities=["episodic", "task_context"],
        )
        slots = {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T],
            InjectionSlot.EPISODIC: [ITEM_OLD],
        }
        result = m.allocate_slots(3, slots)
        # With episodic first: episodic gets 2 tokens, task_context gets 1 (won't fit)
        assert "episodic" in result
        assert "task_context" not in result

    def test_slots_not_in_priority_list_appended(self):
        """Slots missing from slot_priorities are appended in enum order."""
        # Only episodic listed in priorities; task_context not listed → appended last
        m = TokenBudgetManager(
            slot_priorities=["episodic"],
        )
        slots = {
            InjectionSlot.TASK_CONTEXT: [ITEM_2T],
            InjectionSlot.EPISODIC: [ITEM_OLD],
        }
        result = m.allocate_slots(3, slots)
        # episodic listed first → gets 2 tokens first; task_context appended, 1 remaining, won't fit
        assert "episodic" in result
        assert "task_context" not in result


# ---------------------------------------------------------------------------
# compress_to_fit() — compression level selection
# ---------------------------------------------------------------------------

class TestCompressToFit:
    # --- Edge cases ---

    def test_zero_budget_returns_empty(self, manager):
        assert manager.compress_to_fit([ITEM_21T], 0) == []

    def test_negative_budget_returns_empty(self, manager):
        assert manager.compress_to_fit([ITEM_21T], -10) == []

    def test_empty_items_returns_empty(self, manager):
        assert manager.compress_to_fit([], 1000) == []

    # --- Returns correct types ---

    def test_returns_list_of_compressed_items(self, manager):
        result = manager.compress_to_fit([ITEM_2T], 100, CompressionLevel.EVIDENCE)
        assert isinstance(result, list)
        assert all(isinstance(item, CompressedItem) for item in result)

    def test_token_count_matches_content(self, manager):
        result = manager.compress_to_fit([ITEM_21T], 100, CompressionLevel.EVIDENCE)
        assert len(result) == 1
        item = result[0]
        assert item.token_count == manager.count_tokens(item.content)

    def test_original_tokens_recorded(self, manager):
        result = manager.compress_to_fit([ITEM_21T], 100, CompressionLevel.EVIDENCE)
        assert len(result) == 1
        # original_tokens for ITEM_21T = 21
        assert result[0].original_tokens == 21

    def test_compression_ratio_property(self, manager):
        result = manager.compress_to_fit([ITEM_21T], 100, CompressionLevel.EVIDENCE)
        item = result[0]
        expected = item.token_count / item.original_tokens
        assert item.compression_ratio == pytest.approx(expected)

    def test_tokens_saved_property(self, manager):
        result = manager.compress_to_fit([ITEM_21T], 15, CompressionLevel.POINTER)
        item = result[0]
        assert item.tokens_saved == max(0, item.original_tokens - item.token_count)

    def test_tokens_saved_zero_when_no_compression(self, manager):
        """Item at EVIDENCE with original content → tokens_saved == 0."""
        result = manager.compress_to_fit([ITEM_2T], 100, CompressionLevel.EVIDENCE)
        item = result[0]
        assert item.tokens_saved == 0
        assert item.compression_ratio == pytest.approx(1.0)

    # --- EVIDENCE level (min_compression=EVIDENCE, item fits without truncation) ---

    def test_evidence_level_when_item_fits_and_min_evidence(self, manager):
        # Explicitly allow EVIDENCE level
        result = manager.compress_to_fit([ITEM_2T], 100, CompressionLevel.EVIDENCE)
        assert len(result) == 1
        assert result[0].compression == CompressionLevel.EVIDENCE
        assert result[0].content == ITEM_2T

    def test_evidence_level_full_content_preserved(self, manager):
        result = manager.compress_to_fit([ITEM_21T], 100, CompressionLevel.EVIDENCE)
        assert result[0].compression == CompressionLevel.EVIDENCE
        assert result[0].content == ITEM_21T

    # --- COMPACT level (default min_compression=COMPACT) ---

    def test_compact_default_min_blocks_evidence(self, manager):
        """Default min_compression=COMPACT: EVIDENCE is blocked even when item fits.

        Audit fix: min_compression=COMPACT correctly prevents EVIDENCE level.
        Item goes through COMPACT truncation instead, and since 21 tokens >
        pointer_max (12), it gets labeled COMPACT.
        """
        result = manager.compress_to_fit([ITEM_21T], 100)
        assert len(result) == 1
        # Audit fix: EVIDENCE is blocked by default min_compression=COMPACT
        assert result[0].compression == CompressionLevel.COMPACT

    def test_compact_level_when_truncation_required(self, manager):
        """min=COMPACT, item=21 tokens, budget=18: must truncate → COMPACT (18 ≤ 35)."""
        # Item doesn't fit at EVIDENCE (21 > 18), so falls to COMPACT branch.
        # Truncated to 18 tokens; 18 ≤ compact_max(35) → COMPACT.
        result = manager.compress_to_fit([ITEM_21T], 18)
        assert result[0].compression == CompressionLevel.COMPACT
        assert result[0].token_count <= 18

    def test_compact_level_truncation(self, manager):
        """min=COMPACT, item=21 tokens, budget=18 → COMPACT with 18 tokens (truncated)."""
        result = manager.compress_to_fit([ITEM_21T], 18)
        assert len(result) == 1
        item = result[0]
        assert item.compression == CompressionLevel.COMPACT
        assert item.token_count <= 18

    # --- POINTER level ---

    def test_pointer_level_tight_budget_min_compact(self, manager):
        """min=COMPACT, budget=10 (≤ pointer_max=15): compression is POINTER."""
        result = manager.compress_to_fit([ITEM_21T], 10)
        assert result[0].compression == CompressionLevel.POINTER
        assert result[0].token_count <= 15

    def test_pointer_level_explicit_min_pointer(self, manager):
        """min=POINTER allows POINTER label."""
        result = manager.compress_to_fit([ITEM_21T], 10, CompressionLevel.POINTER)
        assert result[0].compression == CompressionLevel.POINTER
        assert result[0].token_count <= 15

    def test_pointer_min_only_pointer_candidates(self, manager):
        """min=POINTER: all items truncated to ≤15 tokens regardless of budget."""
        result = manager.compress_to_fit([ITEM_21T], 500, CompressionLevel.POINTER)
        assert len(result) == 1
        assert result[0].compression == CompressionLevel.POINTER
        assert result[0].token_count <= 15

    # --- min_compression=EVIDENCE fallback ---

    def test_evidence_min_falls_back_to_compact_if_item_too_large(self, manager):
        """min=EVIDENCE, item=21 tokens, budget=18: item can't fit at EVIDENCE (21>18),
        falls back to COMPACT."""
        result = manager.compress_to_fit([ITEM_21T], 18, CompressionLevel.EVIDENCE)
        assert len(result) == 1
        # EVIDENCE not possible (21 > 18), COMPACT tried with cap=18
        assert result[0].compression == CompressionLevel.COMPACT
        assert result[0].token_count <= 18

    def test_evidence_min_item_fits_without_truncation(self, manager):
        """min=EVIDENCE, item=2 tokens, budget=100: item fits at EVIDENCE."""
        result = manager.compress_to_fit([ITEM_2T], 100, CompressionLevel.EVIDENCE)
        assert result[0].compression == CompressionLevel.EVIDENCE

    # --- Multiple items and budget exhaustion ---

    def test_multiple_items_fill_budget_greedy(self, manager):
        """Items are processed greedily; budget exhaustion stops further items."""
        items = [ITEM_2T] * 10
        result = manager.compress_to_fit(items, 6, CompressionLevel.EVIDENCE)
        # 3 items × 2 tokens = 6; 4th item won't fit
        total_tokens = sum(i.token_count for i in result)
        assert total_tokens <= 6
        assert len(result) <= 5  # At most 3 × 2 = 6 tokens → 3 items

    def test_multiple_items_all_fit(self, manager):
        items = [ITEM_2T, ITEM_2T]
        result = manager.compress_to_fit(items, 100, CompressionLevel.EVIDENCE)
        assert len(result) == 2
        assert all(r.compression == CompressionLevel.EVIDENCE for r in result)

    def test_item_that_cannot_fit_is_omitted(self, manager):
        """If an item can't fit even at maximum compression, it's skipped."""
        # Budget=1, item is 2 tokens — even pointer truncation gives ≥1 token ("…")
        # but "…" is 1 token so it fits at POINTER level
        # Use a proper scenario where nothing fits (budget=0):
        result = manager.compress_to_fit([ITEM_21T], 0)
        assert result == []

    def test_total_token_count_within_target(self, manager):
        items = [ITEM_21T] * 5
        target = 50
        result = manager.compress_to_fit(items, target, CompressionLevel.COMPACT)
        total = sum(r.token_count for r in result)
        assert total <= target

    def test_single_large_item_truncated_to_fit(self, manager):
        """Single 21-token item compressed into 10-token budget → fits at POINTER."""
        result = manager.compress_to_fit([ITEM_21T], 10)
        assert len(result) == 1
        assert result[0].token_count <= 10

    # --- Compression levels are within expected token caps ---

    def test_evidence_items_within_200_token_cap(self, manager):
        """EVIDENCE items should not exceed the 200-token cap."""
        result = manager.compress_to_fit([ITEM_21T], 500, CompressionLevel.EVIDENCE)
        for item in result:
            if item.compression == CompressionLevel.EVIDENCE:
                assert item.token_count <= 200

    def test_compact_items_within_80_token_cap(self, manager):
        """COMPACT items should not exceed the 80-token cap (audit fix: was 35)."""
        items = [ITEM_21T] * 3
        result = manager.compress_to_fit(items, 100, CompressionLevel.COMPACT)
        for item in result:
            if item.compression == CompressionLevel.COMPACT:
                assert item.token_count <= 80

    def test_pointer_items_within_12_token_cap(self, manager):
        """POINTER items should not exceed the 12-token cap (audit fix: was 15)."""
        items = [ITEM_21T] * 3
        result = manager.compress_to_fit(items, 100, CompressionLevel.POINTER)
        for item in result:
            if item.compression == CompressionLevel.POINTER:
                assert item.token_count <= 12


# ---------------------------------------------------------------------------
# context_window_budgets() — fraction-based splitting
# ---------------------------------------------------------------------------

class TestContextWindowBudgets:
    def test_default_fractions_sum_to_context_window(self, manager):
        cw = 10_000
        result = manager.context_window_budgets(cw)
        assert result["memory"] + result["user"] + result["reserve"] == cw

    def test_default_fractions_correct_values(self, manager):
        cw = 10_000
        result = manager.context_window_budgets(cw)
        assert result["memory"] == 2500   # 0.25 × 10000
        assert result["user"] == 5500     # 0.55 × 10000
        assert result["reserve"] == 2000  # 10000 - 2500 - 5500

    def test_custom_fractions_via_constructor(self):
        m = TokenBudgetManager(
            memory_fraction=0.30,
            user_fraction=0.50,
            reserve_fraction=0.20,
        )
        cw = 10_000
        result = m.context_window_budgets(cw)
        assert result["memory"] == 3000
        assert result["user"] == 5000
        assert result["reserve"] == 2000
        assert result["memory"] + result["user"] + result["reserve"] == cw

    def test_reserve_always_non_negative(self):
        """Even if fractions sum > 1.0, reserve is clamped to 0."""
        # memory=0.60 + user=0.60 = 1.20 → reserve would be negative → clamped to 0
        m = TokenBudgetManager(memory_fraction=0.60, user_fraction=0.60)
        result = m.context_window_budgets(10_000)
        assert result["reserve"] >= 0

    def test_returns_all_three_keys(self, manager):
        result = manager.context_window_budgets(8192)
        assert set(result.keys()) == {"memory", "user", "reserve"}

    def test_fractions_exposed_as_properties(self):
        m = TokenBudgetManager(
            memory_fraction=0.25,
            user_fraction=0.55,
            reserve_fraction=0.20,
        )
        assert m.memory_fraction == pytest.approx(0.25)
        assert m.user_fraction == pytest.approx(0.55)
        assert m.reserve_fraction == pytest.approx(0.20)


# ---------------------------------------------------------------------------
# Config fields and constructor wiring
# ---------------------------------------------------------------------------

class TestConfigFields:
    def test_default_memory_fraction(self, manager):
        assert manager.memory_fraction == pytest.approx(
            TokenBudgetManager.DEFAULT_MEMORY_FRACTION
        )

    def test_default_user_fraction(self, manager):
        assert manager.user_fraction == pytest.approx(
            TokenBudgetManager.DEFAULT_USER_FRACTION
        )

    def test_default_reserve_fraction(self, manager):
        assert manager.reserve_fraction == pytest.approx(
            TokenBudgetManager.DEFAULT_RESERVE_FRACTION
        )

    def test_constructor_memory_fraction_override(self):
        m = TokenBudgetManager(memory_fraction=0.10)
        assert m.memory_fraction == pytest.approx(0.10)

    def test_constructor_user_fraction_override(self):
        m = TokenBudgetManager(user_fraction=0.70)
        assert m.user_fraction == pytest.approx(0.70)

    def test_constructor_reserve_fraction_override(self):
        m = TokenBudgetManager(reserve_fraction=0.15)
        assert m.reserve_fraction == pytest.approx(0.15)

    def test_config_duck_type_fractions(self):
        """Constructor reads fractions from a config duck-type when available."""
        class FakeConfig:
            total_memory_token_budget = 4000
            cornerstone_token_budget = 800
            token_budget_memory_fraction = 0.28
            token_budget_user_fraction = 0.52
            token_budget_reserve_fraction = 0.20

        m = TokenBudgetManager(config=FakeConfig())
        assert m.memory_fraction == pytest.approx(0.28)
        assert m.user_fraction == pytest.approx(0.52)
        assert m.reserve_fraction == pytest.approx(0.20)

    def test_constructor_override_wins_over_config(self):
        """Explicit constructor args take precedence over config."""
        class FakeConfig:
            total_memory_token_budget = 4000
            cornerstone_token_budget = 800
            token_budget_memory_fraction = 0.10
            token_budget_user_fraction = 0.80
            token_budget_reserve_fraction = 0.10

        m = TokenBudgetManager(config=FakeConfig(), memory_fraction=0.30)
        assert m.memory_fraction == pytest.approx(0.30)

    def test_slot_priorities_constructor(self):
        custom = ["episodic", "procedural", "task_context"]
        m = TokenBudgetManager(slot_priorities=custom)
        assert m.slot_priorities == custom

    def test_default_slot_priorities(self, manager):
        expected = InjectionSlot.default_priority_names()
        assert manager.slot_priorities == expected


# ---------------------------------------------------------------------------
# Backward compatibility: allocate_layers() unchanged
# ---------------------------------------------------------------------------

class TestAllocateLayersBackwardCompat:
    def test_allocate_layers_returns_budget_allocation(self, manager):
        result = manager.allocate_layers({
            MemoryLayer.CORNERSTONE: ["cornerstone text"],
            MemoryLayer.SESSION: ["session item"],
        })
        assert isinstance(result, BudgetAllocation)

    def test_allocate_layers_cornerstones_reserved(self, manager):
        result = manager.allocate_layers({
            MemoryLayer.CORNERSTONE: ["cornerstone memory text"],
        })
        assert result.cornerstone_reserved == manager.cornerstone_budget
        assert "cornerstone" in result.allocations

    def test_allocate_layers_priority_order(self, manager):
        """Session layer has higher priority than procedural; both should fit."""
        result = manager.allocate_layers({
            MemoryLayer.SESSION: ["session item"],
            MemoryLayer.SEMANTIC: ["semantic fact"],
            MemoryLayer.PROCEDURAL: ["procedural step"],
        })
        assert "session" in result.allocations
        assert "semantic" in result.allocations
        assert "procedural" in result.allocations

    def test_allocate_layers_total_budget_respected(self):
        """Total allocated + cornerstones must not exceed total budget."""
        m = TokenBudgetManager(total_budget=50, cornerstone_budget=20)
        result = m.allocate_layers({
            MemoryLayer.CORNERSTONE: ["cornerstone text"],
            MemoryLayer.SESSION: ["session item"],
            MemoryLayer.SEMANTIC: ["semantic memory claim"],
        })
        assert result.used <= 50

    def test_allocate_layers_not_affected_by_slot_methods(self, manager):
        """Running allocate_slots/compress_to_fit does not alter allocate_layers output."""
        # Pre-run slot operations
        manager.allocate_slots(500, {InjectionSlot.TASK_CONTEXT: [ITEM_2T]})
        manager.compress_to_fit([ITEM_21T], 50)

        # allocate_layers should still work identically
        result = manager.allocate_layers({
            MemoryLayer.CORNERSTONE: ["cornerstone memory"],
        })
        assert isinstance(result, BudgetAllocation)
        assert result.total_budget == manager.total_budget

    def test_allocate_layers_total_budget_property(self, manager):
        result = manager.allocate_layers({})
        assert result.total_budget == manager.total_budget

    def test_allocate_layers_with_total_override(self, manager):
        result = manager.allocate_layers({}, total=1000)
        assert result.total_budget == 1000
