github.com/weaviate/weaviate@v1.24.6/usecases/replica/batch_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package replica
    13  
    14  import (
    15  	"sort"
    16  	"strconv"
    17  	"testing"
    18  
    19  	"github.com/go-openapi/strfmt"
    20  	"github.com/stretchr/testify/assert"
    21  	"github.com/weaviate/weaviate/entities/models"
    22  	"github.com/weaviate/weaviate/entities/storobj"
    23  )
    24  
    25  func TestBatchInput(t *testing.T) {
    26  	var (
    27  		N    = 9
    28  		ids  = make([]strfmt.UUID, N)
    29  		data = make([]*storobj.Object, N)
    30  	)
    31  	for i := 0; i < N; i++ {
    32  		uuid := strfmt.UUID(strconv.Itoa(i))
    33  		ids[i] = uuid
    34  		data[i] = objectEx(uuid, 1, "S1", "N1")
    35  	}
    36  	parts := cluster(createBatch(data))
    37  	assert.Len(t, parts, 1)
    38  	assert.Equal(t, parts[0], shardPart{
    39  		Shard: "S1",
    40  		Node:  "N1",
    41  		Data:  data,
    42  		Index: []int{0, 1, 2, 3, 4, 5, 6, 7, 8},
    43  	})
    44  	assert.Equal(t, parts[0].ObjectIDs(), ids)
    45  
    46  	data[0].BelongsToShard = "S2"
    47  	data[0].BelongsToNode = "N2"
    48  	data[2].BelongsToShard = "S2"
    49  	data[2].BelongsToNode = "N2"
    50  	data[3].BelongsToShard = "S2"
    51  	data[4].BelongsToNode = "N2"
    52  	data[5].BelongsToShard = "S2"
    53  	data[5].BelongsToNode = "N2"
    54  
    55  	parts = cluster(createBatch(data))
    56  	sort.Slice(parts, func(i, j int) bool { return len(parts[i].Index) < len(parts[j].Index) })
    57  	assert.Len(t, parts, 2)
    58  	assert.ElementsMatch(t, parts[0].ObjectIDs(), []strfmt.UUID{ids[0], ids[2], ids[3], ids[5]})
    59  	assert.Equal(t, parts[0].Shard, "S2")
    60  	assert.Equal(t, parts[0].Node, "N2")
    61  
    62  	assert.ElementsMatch(t, parts[1].ObjectIDs(), []strfmt.UUID{ids[1], ids[4], ids[6], ids[7], ids[8]})
    63  	assert.Equal(t, parts[1].Shard, "S1")
    64  	assert.Equal(t, parts[1].Node, "N1")
    65  }
    66  
    67  func genInputs(node, shard string, updateTime int64, ids []strfmt.UUID) ([]*storobj.Object, []RepairResponse) {
    68  	xs := make([]*storobj.Object, len(ids))
    69  	digestR := make([]RepairResponse, len(ids))
    70  	for i, id := range ids {
    71  		xs[i] = &storobj.Object{
    72  			Object: models.Object{
    73  				ID:                 id,
    74  				LastUpdateTimeUnix: updateTime,
    75  			},
    76  			BelongsToShard: shard,
    77  			BelongsToNode:  node,
    78  		}
    79  		digestR[i] = RepairResponse{ID: ids[i].String(), UpdateTime: updateTime}
    80  	}
    81  	return xs, digestR
    82  }
    83  
    84  func setObjectsConsistency(xs []*storobj.Object, isConsistent bool) []*storobj.Object {
    85  	want := make([]*storobj.Object, len(xs))
    86  	for i, x := range xs {
    87  		cp := *x
    88  		cp.IsConsistent = isConsistent
    89  		want[i] = &cp
    90  	}
    91  	return want
    92  }
    93  
    94  func objectEx(id strfmt.UUID, lastTime int64, shard, node string) *storobj.Object {
    95  	return &storobj.Object{
    96  		Object: models.Object{
    97  			ID:                 id,
    98  			LastUpdateTimeUnix: lastTime,
    99  		},
   100  		BelongsToShard: shard,
   101  		BelongsToNode:  node,
   102  	}
   103  }