Move tests to state package, address reviewer feedback