TODO: write the tests