Fix existing unit test
Prepare structure for unit tests