Updated test logic