github.com/lino-network/lino@v0.6.11/testutils/dumper.go (about)

     1  package testutils
     2  
     3  import (
     4  	"fmt"
     5  	"io/ioutil"
     6  	"os"
     7  	"reflect"
     8  
     9  	wire "github.com/cosmos/cosmos-sdk/codec"
    10  	sdk "github.com/cosmos/cosmos-sdk/types"
    11  )
    12  
    13  type ValueInterface interface {
    14  }
    15  
    16  type JSONKV struct {
    17  	Prefix string         `json:"prefix"`
    18  	Key    string         `json:"key"`
    19  	Val    ValueInterface `json:"val"`
    20  }
    21  
    22  type JSONState = []JSONKV
    23  
    24  type prefixMatcher struct {
    25  	prefix      []byte
    26  	mkContainer func() interface{}
    27  	raw         bool
    28  }
    29  
    30  type Dumper struct {
    31  	prefixes  map[string]prefixMatcher
    32  	storeCdc  *wire.Codec
    33  	dumperCdc *wire.Codec
    34  	key       sdk.StoreKey
    35  }
    36  
    37  type OptionCodec func(cdc *wire.Codec)
    38  
    39  func prefixToStr(prefix []byte) string {
    40  	return string([]byte{byte(int(prefix[0]) + int('0'))})
    41  }
    42  
    43  func strToPrefx(str string) []byte {
    44  	return []byte{byte(int([]byte(str)[0]) - int('0'))}
    45  }
    46  
    47  func NewDumper(key sdk.StoreKey, storeCdc *wire.Codec, options ...OptionCodec) *Dumper {
    48  	str := ""
    49  	dumperCdc := wire.New()
    50  	wire.RegisterCrypto(dumperCdc)
    51  	dumperCdc.RegisterInterface((*ValueInterface)(nil), nil)
    52  	dumperCdc.RegisterConcrete(str, "str", nil)
    53  	for _, option := range options {
    54  		option(dumperCdc)
    55  	}
    56  	return &Dumper{
    57  		prefixes:  make(map[string]prefixMatcher),
    58  		storeCdc:  storeCdc,
    59  		dumperCdc: dumperCdc,
    60  		key:       key,
    61  	}
    62  }
    63  
    64  func (d *Dumper) RegisterType(t interface{}, name string, subStore []byte) {
    65  	d.prefixes[string(subStore)] = prefixMatcher{
    66  		subStore,
    67  		func() interface{} {
    68  			return reflect.New(reflect.ValueOf(t).Elem().Type()).Interface()
    69  		},
    70  		false}
    71  	d.dumperCdc.RegisterConcrete(t, name, nil)
    72  }
    73  
    74  func (d *Dumper) RegisterRawString(subStore []byte) {
    75  	d.prefixes[string(subStore)] = prefixMatcher{
    76  		prefix: subStore,
    77  		raw:    true,
    78  	}
    79  }
    80  
    81  func (d *Dumper) ToJSON(ctx sdk.Context) []byte {
    82  	store := ctx.KVStore(d.key)
    83  	state := make([]JSONKV, 0)
    84  	itr := sdk.KVStorePrefixIterator(store, nil)
    85  	defer itr.Close()
    86  	for ; itr.Valid(); itr.Next() {
    87  		key := itr.Key()
    88  		val := itr.Value()
    89  		if len(key) == 0 {
    90  			panic("zero length key")
    91  		}
    92  		prefix, ok := d.prefixes[string(key[0])]
    93  		if !ok {
    94  			panic(fmt.Sprintf("unknown substoreprefix: %d", int(key[0])))
    95  		}
    96  		pre := prefixToStr(key)
    97  		var kv JSONKV
    98  		if !prefix.raw {
    99  			container := prefix.mkContainer()
   100  			d.storeCdc.MustUnmarshalBinaryLengthPrefixed(val, container)
   101  			kv = JSONKV{
   102  				Prefix: pre,
   103  				Key:    string(key[1:]),
   104  				Val:    container,
   105  			}
   106  		} else {
   107  			kv = JSONKV{
   108  				Prefix: pre,
   109  				Key:    string(key[1:]),
   110  				Val:    string(val),
   111  			}
   112  		}
   113  		state = append(state, kv)
   114  	}
   115  	bz, err := d.dumperCdc.MarshalJSONIndent(state, "", "  ")
   116  	if err != nil {
   117  		panic(err)
   118  	}
   119  	return bz
   120  }
   121  
   122  func (d *Dumper) DumpToFile(ctx sdk.Context, filepath string) {
   123  	bz := d.ToJSON(ctx)
   124  
   125  	f, err := os.Create(filepath)
   126  	if err != nil {
   127  		panic(err)
   128  	}
   129  	defer f.Close()
   130  	if _, err = f.Write(bz); err != nil {
   131  		panic(err)
   132  	}
   133  	if err = f.Sync(); err != nil {
   134  		panic(err)
   135  	}
   136  }
   137  
   138  func (d *Dumper) LoadFromFile(ctx sdk.Context, filepath string) {
   139  	f, err := os.Open(filepath)
   140  	if err != nil {
   141  		panic(err)
   142  	}
   143  	defer f.Close()
   144  	bz, err := ioutil.ReadAll(f)
   145  	if err != nil {
   146  		panic(err)
   147  	}
   148  	state := make(JSONState, 0)
   149  	d.dumperCdc.MustUnmarshalJSON(bz, &state)
   150  
   151  	store := ctx.KVStore(d.key)
   152  	for _, v := range state {
   153  		pre := strToPrefx(v.Prefix)
   154  		prefix, ok := d.prefixes[string(pre)]
   155  		if !ok {
   156  			panic(fmt.Sprintf("unknown prefix: %v", v.Prefix))
   157  		}
   158  		if prefix.raw {
   159  			store.Set(append(pre, []byte(v.Key)...), []byte(v.Val.(string)))
   160  		} else {
   161  			store.Set(append(pre, []byte(v.Key)...),
   162  				d.storeCdc.MustMarshalBinaryLengthPrefixed(v.Val))
   163  		}
   164  	}
   165  }