Skip to content

Commit

Permalink
add feature engineering tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@guidopetri
guidopetri committed Sep 1, 2024
1 parent 06862ce commit 0e31973
Show file tree
Hide file tree
Showing 2 changed files with 202 additions and 0 deletions.
16 changes: 16 additions & 0 deletions tests/__snapshots__/test_feature_engineering.ambr
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# serializer version: 1
# name: test_clean_chess_df
'{"event_type":{"0":"Rated bullet game"},"game_link":{"0":"https:\\/\\/lichess.org\\/KvnsPlh9"},"date_played":{"0":"2024.01.29"},"round":{"0":"?"},"white":{"0":"Nalajr"},"black":{"0":"thibault"},"result":{"0":"1-0"},"utc_date_played":{"0":"2024.01.29"},"time_played":{"0":"09:44:48"},"white_elo":{"0":"1827"},"black_elo":{"0":"1794"},"white_rating_diff":{"0":"+5"},"black_rating_diff":{"0":"-14"},"chess_variant":{"0":"Standard"},"time_control":{"0":"120+1"},"opening_played":{"0":"B30"},"lichess_opening":{"0":"Sicilian Defense"},"termination":{"0":"Normal"},"evaluations":{"0":[0.15,0.25,0.24]},"eval_depths":{"0":[20,20,20]},"clocks":{"0":["0:02:00","0:02:00","0:02:00"]},"white_berserked":{"0":false},"black_berserked":{"0":false},"queen_exchange":{"0":false},"castling_sides":{"0":{"black":"kingside","white":"kingside"}},"has_promotion":{"0":false},"promotion_count":{"0":{"True":0,"False":0}},"promotions":{"0":{"True":[],"False":[]}},"promotion_count_white":{"0":0},"promotion_count_black":{"0":0},"promotions_white":{"0":""},"promotions_black":{"0":""},"positions":{"0":["rnbqkbnr\\/pppppppp\\/8\\/8\\/4P3\\/8\\/PPPP1PPP\\/RNBQKBNR b KQkq - 0 1","rnbqkbnr\\/pp1ppppp\\/8\\/2p5\\/4P3\\/8\\/PPPP1PPP\\/RNBQKBNR w KQkq - 0 2","rnbqkbnr\\/pp1ppppp\\/8\\/2p5\\/4P3\\/5N2\\/PPPP1PPP\\/RNBQKB1R b KQkq - 1 2"]},"material_by_move":{"0":[{"p":8,"P":8,"k":1,"K":1},{"p":8,"P":8,"k":1,"K":1},{"p":8,"P":8,"k":1,"K":1}]},"moves":{"0":["e4","c5","Nf3"]},"speed":{"0":"bullet"},"status":{"0":"mate"},"black_elo_tentative":{"0":false},"white_elo_tentative":{"0":false}}'
# ---
# name: test_explode_clocks
'{"game_link":{"0":"https:\\/\\/fake-link.com\\/abc","1":"https:\\/\\/fake-link.com\\/abc","2":"https:\\/\\/fake-link.com\\/abc"},"clock":{"0":99,"1":105,"2":93},"half_move":{"0":1,"1":2,"2":3}}'
# ---
# name: test_explode_materials
'{"game_link":{"0":"https:\\/\\/fake-link.com\\/abc","1":"https:\\/\\/fake-link.com\\/abc"},"pawns_black":{"0":0,"1":9},"rooks_black":{"0":1,"1":0},"knights_black":{"0":2,"1":1},"bishops_black":{"0":3,"1":2},"queens_black":{"0":4,"1":3},"pawns_white":{"0":5,"1":4},"rooks_white":{"0":6,"1":5},"knights_white":{"0":7,"1":6},"bishops_white":{"0":8,"1":7},"queens_white":{"0":9,"1":8},"half_move":{"0":1,"1":2}}'
# ---
# name: test_explode_moves
'{"game_link":{"0":"https:\\/\\/fake-link.com\\/abc","1":"https:\\/\\/fake-link.com\\/abc","2":"https:\\/\\/fake-link.com\\/abc"},"move":{"0":"e4","1":"c5","2":"Nf3"},"half_move":{"0":1,"1":2,"2":3}}'
# ---
# name: test_explode_positions
'{"game_link":{"0":"https:\\/\\/fake-link.com\\/abc","1":"https:\\/\\/fake-link.com\\/abc","2":"https:\\/\\/fake-link.com\\/abc"},"position":{"0":"rnbqkbnr\\/pppppppp\\/8\\/8\\/4P3\\/8\\/PPPP1PPP\\/RNBQKBNR b KQkq - 0 1","1":"rnbqkbnr\\/pp1ppppp\\/8\\/2p5\\/4P3\\/8\\/PPPP1PPP\\/RNBQKBNR w KQkq - 0 2","2":"rnbqkbnr\\/pp1ppppp\\/8\\/2p5\\/4P3\\/5N2\\/PPPP1PPP\\/RNBQKB1R b KQkq - 1 2"},"half_move":{"0":1,"1":2,"2":3},"fen":{"0":"rnbqkbnr\\/pppppppp\\/8\\/8\\/4P3\\/8\\/PPPP1PPP\\/RNBQKBNR b KQkq - 0","1":"rnbqkbnr\\/pp1ppppp\\/8\\/2p5\\/4P3\\/8\\/PPPP1PPP\\/RNBQKBNR w KQkq - 0","2":"rnbqkbnr\\/pp1ppppp\\/8\\/2p5\\/4P3\\/5N2\\/PPPP1PPP\\/RNBQKB1R b KQkq - 1"}}'
# ---
186 changes: 186 additions & 0 deletions tests/test_feature_engineering.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,186 @@
#! /usr/bin/env python3


from collections import Counter

import pandas as pd
from feature_engineering import (
clean_chess_df,
explode_clocks,
explode_materials,
explode_moves,
explode_positions,
)


def test_clean_chess_df(snapshot):
pgn_input_df = pd.DataFrame(
[['Rated bullet game',
'https://lichess.org/KvnsPlh9',
'2024.01.29',
'?',
'Nalajr',
'thibault',
'1-0',
'2024.01.29',
'09:44:48',
'1827',
'1794',
'+5',
'-14',
'Standard',
'120+1',
'B30',
'Sicilian Defense',
'Normal',
[0.15, 0.25, 0.24],
[20, 20, 20],
['0:02:00', '0:02:00', '0:02:00'],
False,
False,
False,
{'black': 'kingside', 'white': 'kingside'},
False,
{True: 0, False: 0},
{True: [], False: []},
0,
0,
'',
'',
['rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq - 0 1',
'rnbqkbnr/pp1ppppp/8/2p5/4P3/8/PPPP1PPP/RNBQKBNR w KQkq - 0 2',
'rnbqkbnr/pp1ppppp/8/2p5/4P3/5N2/PPPP1PPP/RNBQKB1R b KQkq - 1 2',
],
[Counter({'p': 8, 'P': 8, 'k': 1, 'K': 1}),
Counter({'p': 8, 'P': 8, 'k': 1, 'K': 1}),
Counter({'p': 8, 'P': 8, 'k': 1, 'K': 1}),
],
['e4', 'c5', 'Nf3'],
]],
columns=['Event',
'Site',
'Date',
'Round',
'White',
'Black',
'Result',
'UTCDate',
'UTCTime',
'WhiteElo',
'BlackElo',
'WhiteRatingDiff',
'BlackRatingDiff',
'Variant',
'TimeControl',
'ECO',
'Opening',
'Termination',
'evaluations',
'eval_depths',
'clocks',
'white_berserked',
'black_berserked',
'queen_exchange',
'castling_sides',
'has_promotion',
'promotion_count',
'promotions',
'promotion_count_white',
'promotion_count_black',
'promotions_white',
'promotions_black',
'positions',
'material_by_move',
'moves',
])
json_input_df = pd.DataFrame(
[['KvnsPlh9',
True,
'standard',
'bullet',
'bullet',
1706521488926,
1706521757537,
'mate',
'pool',
'white',
'Nalajr',
'nalajr',
1827,
5,
'thibault',
'symbols.rainbow-flag',
True,
'thibault',
1794,
-14,
120,
1,
160,
]],
columns=['id',
'rated',
'variant',
'speed',
'perf',
'createdAt',
'lastMoveAt',
'status',
'source',
'winner',
'players_white_user_name',
'players_white_user_id',
'players_white_rating',
'players_white_ratingDiff',
'players_black_user_name',
'players_black_user_flair',
'players_black_user_patron',
'players_black_user_id',
'players_black_rating',
'players_black_ratingDiff',
'clock_initial',
'clock_increment',
'clock_totalTime',
])
df = clean_chess_df(pgn_input_df, json_input_df)
assert df.reset_index(drop=True).to_json() == snapshot


def test_explode_moves(snapshot):
input_df = pd.DataFrame([['https://fake-link.com/abc',
['e4', 'c5', 'Nf3']]],
columns=['game_link', 'moves'])
df = explode_moves(input_df)
assert df.reset_index(drop=True).to_json() == snapshot


def test_explode_clocks(snapshot):
input_df = pd.DataFrame([['https://fake-link.com/abc',
['0:01:39', '0:01:45', '0:01:33']]],
columns=['game_link', 'clocks'])
df = explode_clocks(input_df)
assert df.reset_index(drop=True).to_json() == snapshot


def test_explode_positions(snapshot):
input_df = pd.DataFrame(
[['https://fake-link.com/abc',
['rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b KQkq - 0 1',
'rnbqkbnr/pp1ppppp/8/2p5/4P3/8/PPPP1PPP/RNBQKBNR w KQkq - 0 2',
'rnbqkbnr/pp1ppppp/8/2p5/4P3/5N2/PPPP1PPP/RNBQKB1R b KQkq - 1 2',
],
]],
columns=['game_link', 'positions'])
df = explode_positions(input_df)
assert df.reset_index(drop=True).to_json() == snapshot


def test_explode_materials(snapshot):
input_df = pd.DataFrame([['https://fake-link.com/abc',
[{k: idx for idx, k in enumerate('prnbqPRNBQ')},
{k: idx for idx, k in enumerate('rnbqPRNBQp')},
]
]],
columns=['game_link', 'material_by_move'])
df = explode_materials(input_df)
assert df.reset_index(drop=True).to_json() == snapshot

0 comments on commit 0e31973

Please sign in to comment.