Fix docstrings and add tests for state store methods