github.com/0chain/gosdk@v1.17.11/zboxcore/sdk/downloadworker_test.go (about)

     1  package sdk
     2  
     3  import (
     4  	"encoding/hex"
     5  	"math/rand"
     6  	"sync"
     7  	"testing"
     8  
     9  	"github.com/0chain/gosdk/zboxcore/zboxutil"
    10  	"github.com/stretchr/testify/require"
    11  	"golang.org/x/crypto/sha3"
    12  )
    13  
    14  func TestRemoveFromMask(t *testing.T) {
    15  	req := DownloadRequest{}
    16  	req.maskMu = &sync.Mutex{}
    17  	N := 30
    18  	req.downloadMask = zboxutil.NewUint128(1).Lsh(uint64(30)).Sub64(1)
    19  
    20  	require.Equal(t, N, req.downloadMask.CountOnes())
    21  
    22  	n := 10
    23  	for i := 0; i < n; i++ {
    24  		req.removeFromMask(uint64(i))
    25  	}
    26  
    27  	expected := N - n
    28  	require.Equal(t, expected, req.downloadMask.CountOnes())
    29  }
    30  
    31  func TestDecodeEC(t *testing.T) {
    32  	type input struct {
    33  		name        string
    34  		req         *DownloadRequest
    35  		shards      [][]byte
    36  		wantErr     bool
    37  		contentHash string
    38  		errMsg      string
    39  		setup       func(in *input)
    40  	}
    41  
    42  	tests := []*input{
    43  		{
    44  			name:    "should be ok",
    45  			wantErr: false,
    46  			setup: func(in *input) {
    47  				req := DownloadRequest{}
    48  				req.datashards = 4
    49  				req.parityshards = 2
    50  				req.effectiveBlockSize = 64 * 1024
    51  
    52  				err := req.initEC()
    53  				require.NoError(t, err)
    54  
    55  				d, err := getDummyData(64 * 1024 * 4)
    56  				require.NoError(t, err)
    57  
    58  				h := sha3.New256()
    59  				n, err := h.Write(d)
    60  				require.NoError(t, err)
    61  				require.Equal(t, len(d), n)
    62  				in.contentHash = hex.EncodeToString(h.Sum(nil))
    63  
    64  				shards, err := req.ecEncoder.Split(d)
    65  				require.NoError(t, err)
    66  
    67  				err = req.ecEncoder.Encode(shards)
    68  				require.NoError(t, err)
    69  
    70  				shards[0] = nil
    71  
    72  				in.shards = shards
    73  				in.req = &req
    74  			},
    75  		},
    76  	}
    77  
    78  	for _, test := range tests {
    79  
    80  		test := test
    81  		t.Run(test.name, func(t *testing.T) {
    82  			if test.setup != nil {
    83  				test.setup(test)
    84  			}
    85  
    86  			err := test.req.decodeEC(test.shards)
    87  			if test.wantErr {
    88  				require.Error(t, err)
    89  				require.Contains(t, err.Error(), test.errMsg)
    90  				return
    91  			} else {
    92  				require.NoError(t, err)
    93  			}
    94  
    95  			h := sha3.New256()
    96  			total := 0
    97  			for i := 0; i < test.req.datashards; i++ {
    98  				n, err := h.Write(test.shards[i])
    99  				require.NoError(t, err)
   100  				total += n
   101  			}
   102  			hash := hex.EncodeToString(h.Sum(nil))
   103  			require.Equal(t, test.contentHash, hash)
   104  
   105  		})
   106  	}
   107  }
   108  
   109  func TestFillShards(t *testing.T) {
   110  	type input struct {
   111  		name          string
   112  		wantErr       bool
   113  		totalBlocks   int
   114  		totalBlobbers int
   115  		blobberIdx    int
   116  		expectedSize  int
   117  		req           *DownloadRequest
   118  		setup         func(in *input)
   119  		shards        [][][]byte
   120  		result        *downloadBlock
   121  	}
   122  
   123  	tests := []*input{
   124  		{
   125  			name:    "fill shards ok",
   126  			wantErr: false,
   127  			setup: func(in *input) {
   128  				in.expectedSize = 64 * 1024
   129  				in.totalBlobbers = 4
   130  				in.totalBlocks = 2
   131  				in.blobberIdx = 1
   132  				d, err := getDummyData(in.expectedSize * in.totalBlocks)
   133  				require.NoError(t, err)
   134  				in.req = &DownloadRequest{}
   135  				in.req.maskMu = &sync.Mutex{}
   136  				shards := make([][]byte, in.totalBlocks)
   137  				for i := 0; i < in.totalBlocks; i++ {
   138  					index := i * in.expectedSize
   139  					shards[i] = d[index : index+in.expectedSize]
   140  				}
   141  
   142  				in.result = &downloadBlock{
   143  					BlockChunks: shards,
   144  					Success:     true,
   145  					idx:         in.blobberIdx,
   146  				}
   147  
   148  				in.shards = make([][][]byte, in.totalBlocks)
   149  				for i := range in.shards {
   150  					in.shards[i] = make([][]byte, in.totalBlobbers)
   151  				}
   152  			},
   153  		},
   154  	}
   155  
   156  	for _, test := range tests {
   157  		t.Run(test.name, func(t *testing.T) {
   158  			if test.setup != nil {
   159  				test.setup(test)
   160  			}
   161  
   162  			// maskCount := test.req.downloadMask.CountOnes()
   163  			err := test.req.fillShards(test.shards, test.result)
   164  			if test.wantErr {
   165  				require.Error(t, err)
   166  				// require.Equal(t, maskCount-1, test.req.downloadMask.CountOnes())
   167  				return
   168  			}
   169  
   170  			require.NoError(t, err)
   171  			for i := 0; i < test.totalBlocks; i++ {
   172  				data := test.shards[i][test.blobberIdx]
   173  				require.Equal(t, test.expectedSize, len(data))
   174  			}
   175  
   176  		})
   177  	}
   178  }
   179  
   180  func getDummyData(size int) ([]byte, error) {
   181  	b := make([]byte, size)
   182  	_, err := rand.Read(b) //nolint
   183  	if err != nil {
   184  		return nil, err
   185  	}
   186  	return b, nil
   187  }