test_codec.py

1import unittest
2import codec as c
3chess=c.chess
4class Codec(unittest.TestCase):
5    def test_layout_shift_and_fen_only(self):
6        b=chess.Board();m=chess.Move.from_uci('e2e4');e=c.example(b,m)
7        self.assertEqual((len(e['input_ids']),len(e['labels'])),(122,3));self.assertEqual(e['labels'],[12,28,64])
8        self.assertEqual(e['input_ids'][119:],[c.NEXT,12,28]);self.assertEqual(sum(e['loss_mask']),3)
9        self.assertEqual(e['attention_mask'][2:50],[0]*48);self.assertEqual(e['input_ids'][50],c.STATE)
10        self.assertEqual(e['input_ids'][51:59],[132,130,131,133,134,131,130,132]);self.assertNotIn(c.EOS,e['input_ids'])
11    def test_history_fixed_positions_and_parity(self):
12        b=chess.Board()
13        for u in ['e2e4','e7e5','g1f3','b8c6']:b.push_uci(u)
14        e=c.example(b,chess.Move.from_uci('f1b5'));p=c.prompt(b)
15        self.assertEqual(e['input_ids'][:120],p['input_ids']);self.assertEqual(p['position_ids'],list(range(120)))
16        self.assertEqual(sum(p['attention_mask'][2:50]),12)
17        dropped=c.prompt(b,0);self.assertEqual(dropped['input_ids'][50:],p['input_ids'][50:]);self.assertEqual(dropped['input_ids'][1],c.HIST_NONE)
18        replay=c.from_history(chess.STARTING_FEN,['e2e4','e7e5','g1f3','b8c6'])
19        self.assertEqual(c.prompt(replay),p)
20    def test_castling_and_ep(self):
21        b=chess.Board('r3k2r/8/8/8/8/8/8/R3K2R w KQkq - 0 1')
22        self.assertEqual(c.prompt(b)['input_ids'][116],c.CASTLE_BASE+15)
23        for u in ('e1g1','e1c1'):self.assertEqual(c.decode(c.example(b,chess.Move.from_uci(u))['labels']).uci(),u)
24        b=chess.Board();b.push_uci('e2e4');self.assertEqual(c.prompt(b)['input_ids'][117],c.EP_BASE+chess.E3)
25        other=chess.Board(b.fen());other.ep_square=None;self.assertEqual(c.position_key(b),c.position_key(other))
26        b=chess.Board('4k3/8/8/3pP3/8/8/8/4K3 w - d6 0 2')
27        self.assertEqual(c.example(b,chess.Move.from_uci('e5d6'))['labels'],[36,43,64])
28        other=b.copy();other.ep_square=None;self.assertNotEqual(c.position_key(b),c.position_key(other))
29    def test_promotions_and_capture_underpromotion(self):
30        for suffix in 'nbrq':
31            for u in ('a7a8'+suffix,'a7b8'+suffix):
32                b=chess.Board('1r2k3/P7/8/8/8/8/8/4K3 w - - 0 1');m=chess.Move.from_uci(u)
33                self.assertEqual(c.decode(c.example(b,m)['labels']),m)
34    def test_clock_and_terminal(self):
35        b=chess.Board('4k2r/8/8/8/8/8/8/R3K3 w - - 149 237')
36        self.assertEqual(c.prompt(b)['input_ids'][118],c.HALF_BASE+149);self.assertEqual(c.prompt(b)['metadata']['fullmove_number'],237)
37        b.halfmove_clock=150;self.assertEqual(c.prompt(b)['input_ids'][118],c.HALF_GE150)
38        with self.assertRaises(ValueError):c.example(b,chess.Move.from_uci('a1a2'))
39    def test_reserved_phase_and_illegal(self):
40        b=chess.Board()
41        for t in ([12,28,72],[12,12,64],[0,1],[0,1,512]):
42            with self.assertRaises(ValueError):c.decode(t)
43        with self.assertRaises(ValueError):c.example(b,chess.Move.from_uci('e2e5'))
44        with self.assertRaises(ValueError):c.masks(b,(64,))
45        self.assertFalse(any(c.masks(b,legal=True)[128:]));self.assertFalse(c.masks(b,legal=True)[c.EOS])
46    def test_variants_and_corrupt_history(self):
47        import chess.variant
48        with self.assertRaises(ValueError):c.prompt(chess.variant.AtomicBoard())
49        with self.assertRaises(ValueError):c.from_history(chess.STARTING_FEN,['e2e5'])
50        b=chess.Board();b.push_uci('e2e4');b.halfmove_clock=99
51        with self.assertRaises(ValueError):c.prompt(b)
52    def test_position_keys_ignore_clocks_not_rights_turn(self):
53        b=chess.Board();o=b.copy();o.halfmove_clock=7;o.fullmove_number=22
54        self.assertEqual(c.position_key(b),c.position_key(o));o.turn=False
55        self.assertNotEqual(c.position_key(b),c.position_key(o));o=b.copy();o.castling_rights=0
56        self.assertNotEqual(c.position_key(b),c.position_key(o))
57if __name__=='__main__':unittest.main()
58