Add tests for energy calculation