Migration script rewrite with unit tests