Some test framework additions

Signed-off-by: Leonard Lyubich <ctulhurider@gmail.com>

balance/migration_test.go

@@ -1,4 +1,4 @@
-package balance
+package balance_test
 
 import (
 	"testing"

tests/migration/model.go

@@ -1,6 +1,7 @@
 package migration
 
 import (
+	"bytes"
 	"testing"
 )
 
@@ -91,6 +92,33 @@ func PersistedItemsWithKeyLength(ln int) ItemGroupModel {
 	}
 }
 
+// PersistedItemsWithPrefix describes persistence of the storage items with
+// specified key prefix.
+func PersistedItemsWithPrefix(prefix string) ItemGroupModel {
+	return ItemGroupModel{
+		keyPattern: func(key []byte) bool {
+			return bytes.HasPrefix(key, []byte(prefix))
+		},
+		modifier: persistedModifier,
+	}
+}
+
+// ReplacedItemsWithKeyLength describes replacement of the storage items with
+// fixed-length keys.
+func ReplacedItemsWithKeyLength(ln int, changeKey func(prevKey []byte) []byte) ItemGroupModel {
+	return ItemGroupModel{
+		keyPattern: func(key []byte) bool {
+			return len(key) == ln
+		},
+		modifier: func(prev kv) kv {
+			return kv{
+				k: changeKey(prev.k),
+				v: prev.v,
+			}
+		},
+	}
+}
+
 // apply applies the UpdateModel to pre-update storage items of some contract
 // and returns list of expected post-update items. Fails if any of items is not
 // covered by the UpdateModel.
@@ -103,7 +131,7 @@ nextItem:
 				continue nextItem
 			}
 		}
-		tb.Fatalf("unexpected storage item '%s', did you forget to forget to include it in the model?")
+		tb.Fatalf("unexpected storage item '%s', did you forget to forget to include it in the model?", in[i].k)
 	}
 	return
 }