Fix unit tests for new json structure