github.com/batchcorp/thrift-iterator@v0.0.0-20220918180557-4c4a158fc6e9/test/level_1/map_test.go (about)

     1  package test
     2  
     3  import (
     4  	"github.com/apache/thrift/lib/go/thrift"
     5  	"github.com/batchcorp/thrift-iterator/general"
     6  	"github.com/batchcorp/thrift-iterator/protocol"
     7  	"github.com/batchcorp/thrift-iterator/raw"
     8  	"github.com/batchcorp/thrift-iterator/test"
     9  	"github.com/stretchr/testify/require"
    10  	"testing"
    11  )
    12  
    13  func Test_decode_map_by_iterator(t *testing.T) {
    14  	should := require.New(t)
    15  	for _, c := range test.Combinations {
    16  		buf, proto := c.CreateProtocol()
    17  		proto.WriteMapBegin(thrift.STRING, thrift.I64, 3)
    18  		proto.WriteString("k1")
    19  		proto.WriteI64(1)
    20  		proto.WriteString("k2")
    21  		proto.WriteI64(2)
    22  		proto.WriteString("k3")
    23  		proto.WriteI64(3)
    24  		proto.WriteMapEnd()
    25  		iter := c.CreateIterator(buf.Bytes())
    26  		keyType, elemType, length := iter.ReadMapHeader()
    27  		should.Equal(protocol.TypeString, keyType)
    28  		should.Equal(protocol.TypeI64, elemType)
    29  		should.Equal(3, length)
    30  		should.Equal("k1", iter.ReadString())
    31  		should.Equal(uint64(1), iter.ReadUint64())
    32  		should.Equal("k2", iter.ReadString())
    33  		should.Equal(uint64(2), iter.ReadUint64())
    34  		should.Equal("k3", iter.ReadString())
    35  		should.Equal(uint64(3), iter.ReadUint64())
    36  	}
    37  }
    38  
    39  func Test_encode_map_by_stream(t *testing.T) {
    40  	should := require.New(t)
    41  	for _, c := range test.Combinations {
    42  		stream := c.CreateStream()
    43  		stream.WriteMapHeader(protocol.TypeString, protocol.TypeI64, 3)
    44  		stream.WriteString("k1")
    45  		stream.WriteUint64(1)
    46  		stream.WriteString("k2")
    47  		stream.WriteUint64(2)
    48  		stream.WriteString("k3")
    49  		stream.WriteUint64(3)
    50  		iter := c.CreateIterator(stream.Buffer())
    51  		keyType, elemType, length := iter.ReadMapHeader()
    52  		should.Equal(protocol.TypeString, keyType)
    53  		should.Equal(protocol.TypeI64, elemType)
    54  		should.Equal(3, length)
    55  		should.Equal("k1", iter.ReadString())
    56  		should.Equal(uint64(1), iter.ReadUint64())
    57  		should.Equal("k2", iter.ReadString())
    58  		should.Equal(uint64(2), iter.ReadUint64())
    59  		should.Equal("k3", iter.ReadString())
    60  		should.Equal(uint64(3), iter.ReadUint64())
    61  	}
    62  }
    63  
    64  func Test_skip_map(t *testing.T) {
    65  	should := require.New(t)
    66  	for _, c := range test.Combinations {
    67  		buf, proto := c.CreateProtocol()
    68  		proto.WriteMapBegin(thrift.I32, thrift.I64, 3)
    69  		proto.WriteI32(1)
    70  		proto.WriteI64(1)
    71  		proto.WriteI32(2)
    72  		proto.WriteI64(2)
    73  		proto.WriteI32(3)
    74  		proto.WriteI64(3)
    75  		proto.WriteMapEnd()
    76  		iter := c.CreateIterator(buf.Bytes())
    77  		should.Equal(buf.Bytes(), iter.SkipMap(nil))
    78  	}
    79  }
    80  
    81  func Test_unmarshal_general_map(t *testing.T) {
    82  	should := require.New(t)
    83  	for _, c := range test.Combinations {
    84  		buf, proto := c.CreateProtocol()
    85  		proto.WriteMapBegin(thrift.I32, thrift.I64, 3)
    86  		proto.WriteI32(1)
    87  		proto.WriteI64(1)
    88  		proto.WriteI32(2)
    89  		proto.WriteI64(2)
    90  		proto.WriteI32(3)
    91  		proto.WriteI64(3)
    92  		proto.WriteMapEnd()
    93  		var val general.Map
    94  		should.NoError(c.Unmarshal(buf.Bytes(), &val))
    95  		should.Equal(general.Map{
    96  			int32(1): int64(1),
    97  			int32(2): int64(2),
    98  			int32(3): int64(3),
    99  		}, val)
   100  	}
   101  }
   102  
   103  func Test_unmarshal_raw_map(t *testing.T) {
   104  	should := require.New(t)
   105  	for _, c := range test.Combinations {
   106  		buf, proto := c.CreateProtocol()
   107  		proto.WriteMapBegin(thrift.I32, thrift.I64, 3)
   108  		proto.WriteI32(1)
   109  		proto.WriteI64(1)
   110  		proto.WriteI32(2)
   111  		proto.WriteI64(2)
   112  		proto.WriteI32(3)
   113  		proto.WriteI64(3)
   114  		proto.WriteMapEnd()
   115  		var val raw.Map
   116  		should.NoError(c.Unmarshal(buf.Bytes(), &val))
   117  		should.Equal(3, len(val.Entries))
   118  		should.Equal(protocol.TypeI32, val.KeyType)
   119  		should.Equal(protocol.TypeI64, val.ElementType)
   120  		iter := c.CreateIterator(val.Entries[int32(1)].Element)
   121  		should.Equal(int64(1), iter.ReadInt64())
   122  	}
   123  }
   124  
   125  func Test_unmarshal_map(t *testing.T) {
   126  	should := require.New(t)
   127  	for _, c := range test.UnmarshalCombinations {
   128  		buf, proto := c.CreateProtocol()
   129  		proto.WriteMapBegin(thrift.I32, thrift.I64, 3)
   130  		proto.WriteI32(1)
   131  		proto.WriteI64(1)
   132  		proto.WriteI32(2)
   133  		proto.WriteI64(2)
   134  		proto.WriteI32(3)
   135  		proto.WriteI64(3)
   136  		proto.WriteMapEnd()
   137  		val := map[int32]int64{}
   138  		should.NoError(c.Unmarshal(buf.Bytes(), &val))
   139  		should.Equal(map[int32]int64{
   140  			int32(1): int64(1),
   141  			int32(2): int64(2),
   142  			int32(3): int64(3),
   143  		}, val)
   144  	}
   145  }
   146  
   147  func Test_marshal_general_map(t *testing.T) {
   148  	should := require.New(t)
   149  	for _, c := range test.Combinations {
   150  		m := general.Map{
   151  			int32(1): int64(1),
   152  			int32(2): int64(2),
   153  			int32(3): int64(3),
   154  		}
   155  
   156  		output, err := c.Marshal(m)
   157  		should.NoError(err)
   158  		output1, err := c.Marshal(&m)
   159  		should.NoError(err)
   160  		var val, val1 general.Map
   161  		should.NoError(c.Unmarshal(output, &val))
   162  		should.NoError(c.Unmarshal(output1, &val1))
   163  		should.Equal(val, val1)
   164  		should.Equal(general.Map{
   165  			int32(1): int64(1),
   166  			int32(2): int64(2),
   167  			int32(3): int64(3),
   168  		}, val)
   169  	}
   170  }
   171  
   172  func Test_marshal_raw_map(t *testing.T) {
   173  	should := require.New(t)
   174  	for _, c := range test.Combinations {
   175  		buf, proto := c.CreateProtocol()
   176  		proto.WriteMapBegin(thrift.I32, thrift.I64, 3)
   177  		proto.WriteI32(1)
   178  		proto.WriteI64(1)
   179  		proto.WriteI32(2)
   180  		proto.WriteI64(2)
   181  		proto.WriteI32(3)
   182  		proto.WriteI64(3)
   183  		proto.WriteMapEnd()
   184  		var val raw.Map
   185  		should.NoError(c.Unmarshal(buf.Bytes(), &val))
   186  
   187  		output, err := c.Marshal(val)
   188  		should.NoError(err)
   189  		output1, err := c.Marshal(&val)
   190  		should.NoError(err)
   191  		var generalVal, generalVal1 general.Map
   192  		should.NoError(c.Unmarshal(output, &generalVal))
   193  		should.NoError(c.Unmarshal(output1, &generalVal1))
   194  		should.Equal(generalVal, generalVal1)
   195  		should.Equal(general.Map{
   196  			int32(1): int64(1),
   197  			int32(2): int64(2),
   198  			int32(3): int64(3),
   199  		}, generalVal)
   200  	}
   201  }
   202  
   203  func Test_marshal_map(t *testing.T) {
   204  	should := require.New(t)
   205  	for _, c := range test.MarshalCombinations {
   206  		m := map[string]int64{
   207  			"k1": int64(1),
   208  			"k2": int64(2),
   209  			"k3": int64(3),
   210  		}
   211  
   212  		output, err := c.Marshal(m)
   213  		should.NoError(err)
   214  		output1, err := c.Marshal(&m)
   215  		should.NoError(err)
   216  		var val, val1 general.Map
   217  		should.NoError(c.Unmarshal(output, &val))
   218  		should.NoError(c.Unmarshal(output1, &val1))
   219  		should.Equal(val, val1)
   220  		should.Equal(general.Map{
   221  			"k1": int64(1),
   222  			"k2": int64(2),
   223  			"k3": int64(3),
   224  		}, val)
   225  	}
   226  }
   227  
   228  func Test_marshal_empty_map(t *testing.T) {
   229  	should := require.New(t)
   230  	for _, c := range test.MarshalCombinations {
   231  		output, err := c.Marshal(map[string]int64{})
   232  		should.NoError(err)
   233  		var val general.Map
   234  		should.NoError(c.Unmarshal(output, &val))
   235  		should.Equal(general.Map{}, val)
   236  	}
   237  }