github.com/weaviate/weaviate@v1.24.6/usecases/sharding/remote_index_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 sharding
    13  
    14  import (
    15  	"context"
    16  	"errors"
    17  	"fmt"
    18  	"testing"
    19  )
    20  
    21  var errAny = errors.New("anyErr")
    22  
    23  func TestQueryReplica(t *testing.T) {
    24  	var (
    25  		ctx                      = context.Background()
    26  		canceledCtx, cancledFunc = context.WithCancel(ctx)
    27  	)
    28  	cancledFunc()
    29  	doIf := func(targetNode string) func(node, host string) (interface{}, error) {
    30  		return func(node, host string) (interface{}, error) {
    31  			if node != targetNode {
    32  				return nil, errAny
    33  			}
    34  			return node, nil
    35  		}
    36  	}
    37  	tests := []struct {
    38  		ctx        context.Context
    39  		resolver   fakeNodeResolver
    40  		schema     fakeSchema
    41  		targetNode string
    42  		success    bool
    43  		name       string
    44  	}{
    45  		{
    46  			ctx, newFakeResolver(0, 0), newFakeSchema(0, 0), "N0", false, "empty schema",
    47  		},
    48  		{
    49  			ctx, newFakeResolver(0, 1), newFakeSchema(1, 2), "N2", false, "unresolved name",
    50  		},
    51  		{
    52  			ctx, newFakeResolver(0, 1), newFakeSchema(0, 1), "N0", true, "one replica",
    53  		},
    54  		{
    55  			ctx, newFakeResolver(0, 9), newFakeSchema(0, 9), "N2", true, "random selection",
    56  		},
    57  		{
    58  			canceledCtx, newFakeResolver(0, 9), newFakeSchema(0, 9), "N2", false, "canceled",
    59  		},
    60  	}
    61  
    62  	for _, test := range tests {
    63  		rindex := RemoteIndex{"C", &test.schema, nil, &test.resolver}
    64  		got, lastNode, err := rindex.queryReplicas(test.ctx, "S", doIf(test.targetNode))
    65  		if !test.success {
    66  			if got != nil {
    67  				t.Errorf("%s: want: nil, got: %v", test.name, got)
    68  			} else if err == nil {
    69  				t.Errorf("%s: must return an error", test.name)
    70  			}
    71  			continue
    72  		}
    73  		if lastNode != test.targetNode {
    74  			t.Errorf("%s: last responding node want:%s got:%s", test.name, test.targetNode, lastNode)
    75  		}
    76  	}
    77  }
    78  
    79  func newFakeResolver(fromNode, toNode int) fakeNodeResolver {
    80  	m := make(map[string]string, toNode-fromNode)
    81  	for i := fromNode; i < toNode; i++ {
    82  		m[fmt.Sprintf("N%d", i)] = fmt.Sprintf("H%d", i)
    83  	}
    84  	return fakeNodeResolver{m}
    85  }
    86  
    87  func newFakeSchema(fromNode, toNode int) fakeSchema {
    88  	nodes := make([]string, 0, toNode-fromNode)
    89  	for i := fromNode; i < toNode; i++ {
    90  		nodes = append(nodes, fmt.Sprintf("N%d", i))
    91  	}
    92  	return fakeSchema{nodes}
    93  }
    94  
    95  type fakeNodeResolver struct {
    96  	rTable map[string]string
    97  }
    98  
    99  func (f *fakeNodeResolver) NodeHostname(name string) (string, bool) {
   100  	host, ok := f.rTable[name]
   101  	return host, ok
   102  }
   103  
   104  type fakeSchema struct {
   105  	nodes []string
   106  }
   107  
   108  func (f *fakeSchema) ShardOwner(class, shard string) (string, error) {
   109  	return "", nil
   110  }
   111  
   112  func (f *fakeSchema) ShardReplicas(class, shard string) ([]string, error) {
   113  	return f.nodes, nil
   114  }