update tests to handle new return data from states