[WIP] Improved tests