Fix the InMemoryShardStateIO used in ShardStateIntegrationTest