import random
import unittest
import torch
from beamboot.rope_memory import RotaryCacheMemory


class RotaryMemoryTests(unittest.TestCase):
    def check_memory(self, dtype):
        memory = RotaryCacheMemory(dtype)
        reference = [0]*16
        rng = random.Random(112)
        for _ in range(200):
            address, value = rng.randrange(16), rng.randrange(16)
            memory.write(address, value)
            reference[address] = value
            self.assertEqual(memory.read(address), value)
        for address, value in enumerate(reference):
            self.assertEqual(memory.read(address), value)
        self.assertEqual(memory.position, 432)

    def test_float32(self): self.check_memory(torch.float32)
    def test_bfloat16(self): self.check_memory(torch.bfloat16)

    def test_evicted_write_is_not_reliably_readable(self):
        memory = RotaryCacheMemory(window=32)
        memory.write(0, 7)
        for _ in range(32): memory.write(1, 3)
        self.assertNotEqual(memory.read(0), 7)
        self.assertEqual(memory.read(1), 3)