import unittest
from src.aes import s_box

class TestAES(unittest.TestCase):
    def test_sbox_known_values(self):
        """Test S-box with known AES standard values."""
        self.assertEqual(s_box(0x00), 0x63)
        self.assertEqual(s_box(0x01), 0x7c)
        self.assertEqual(s_box(0x10), 0xca)
        self.assertEqual(s_box(0x53), 0xed)
        self.assertEqual(s_box(0xff), 0x16)
    
    def test_sbox_invalid_values(self):
        """Test S-box with invalid input values."""
        # Test values outside valid range
        with self.assertRaises(IndexError):
            s_box(-1)
        
        with self.assertRaises(IndexError):
            s_box(256)
        
        with self.assertRaises(IndexError):
            s_box(1000)

if __name__ == '__main__':
    unittest.main()