fix unit tests

@infynyxx review?