Add unit tests for future implementation