github.com/weaviate/weaviate@v1.24.6/usecases/replica/mocks_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  	"context"
    16  	"fmt"
    17  
    18  	"github.com/go-openapi/strfmt"
    19  	"github.com/stretchr/testify/mock"
    20  	"github.com/weaviate/weaviate/entities/additional"
    21  	"github.com/weaviate/weaviate/entities/search"
    22  	"github.com/weaviate/weaviate/entities/storobj"
    23  	"github.com/weaviate/weaviate/usecases/objects"
    24  )
    25  
    26  type fakeRClient struct {
    27  	mock.Mock
    28  }
    29  
    30  func (f *fakeRClient) FetchObject(ctx context.Context, host, index, shard string,
    31  	id strfmt.UUID, props search.SelectProperties,
    32  	additional additional.Properties,
    33  ) (objects.Replica, error) {
    34  	args := f.Called(ctx, host, index, shard, id, props, additional)
    35  	return args.Get(0).(objects.Replica), args.Error(1)
    36  }
    37  
    38  func (f *fakeRClient) FetchObjects(ctx context.Context, host, index,
    39  	shard string, ids []strfmt.UUID,
    40  ) ([]objects.Replica, error) {
    41  	args := f.Called(ctx, host, index, shard, ids)
    42  	return args.Get(0).([]objects.Replica), args.Error(1)
    43  }
    44  
    45  func (f *fakeRClient) OverwriteObjects(ctx context.Context, host, index, shard string,
    46  	xs []*objects.VObject,
    47  ) ([]RepairResponse, error) {
    48  	args := f.Called(ctx, host, index, shard, xs)
    49  	return args.Get(0).([]RepairResponse), args.Error(1)
    50  }
    51  
    52  func (f *fakeRClient) DigestObjects(ctx context.Context, host, index, shard string,
    53  	ids []strfmt.UUID,
    54  ) ([]RepairResponse, error) {
    55  	args := f.Called(ctx, host, index, shard, ids)
    56  	return args.Get(0).([]RepairResponse), args.Error(1)
    57  }
    58  
    59  type fakeClient struct {
    60  	mock.Mock
    61  }
    62  
    63  func (f *fakeClient) PutObject(ctx context.Context, host, index, shard, requestID string,
    64  	obj *storobj.Object,
    65  ) (SimpleResponse, error) {
    66  	args := f.Called(ctx, host, index, shard, requestID, obj)
    67  	return args.Get(0).(SimpleResponse), args.Error(1)
    68  }
    69  
    70  func (f *fakeClient) DeleteObject(ctx context.Context, host, index, shard, requestID string,
    71  	id strfmt.UUID,
    72  ) (SimpleResponse, error) {
    73  	args := f.Called(ctx, host, index, shard, requestID, id)
    74  	return args.Get(0).(SimpleResponse), args.Error(1)
    75  }
    76  
    77  func (f *fakeClient) MergeObject(ctx context.Context, host, index, shard, requestID string,
    78  	doc *objects.MergeDocument,
    79  ) (SimpleResponse, error) {
    80  	args := f.Called(ctx, host, index, shard, requestID, doc)
    81  	return args.Get(0).(SimpleResponse), args.Error(1)
    82  }
    83  
    84  func (f *fakeClient) PutObjects(ctx context.Context, host, index, shard, requestID string,
    85  	objs []*storobj.Object,
    86  ) (SimpleResponse, error) {
    87  	args := f.Called(ctx, host, index, shard, requestID, objs)
    88  	return args.Get(0).(SimpleResponse), args.Error(1)
    89  }
    90  
    91  func (f *fakeClient) DeleteObjects(ctx context.Context, host, index, shard, requestID string,
    92  	uuids []strfmt.UUID, dryRun bool,
    93  ) (SimpleResponse, error) {
    94  	args := f.Called(ctx, host, index, shard, requestID, uuids, dryRun)
    95  	return args.Get(0).(SimpleResponse), args.Error(1)
    96  }
    97  
    98  func (f *fakeClient) AddReferences(ctx context.Context, host, index, shard, requestID string,
    99  	refs []objects.BatchReference,
   100  ) (SimpleResponse, error) {
   101  	args := f.Called(ctx, host, index, shard, requestID, refs)
   102  	return args.Get(0).(SimpleResponse), args.Error(1)
   103  }
   104  
   105  func (f *fakeClient) Commit(ctx context.Context, host, index, shard, requestID string, resp interface{}) error {
   106  	args := f.Called(ctx, host, index, shard, requestID, resp)
   107  	return args.Error(0)
   108  }
   109  
   110  func (f *fakeClient) Abort(ctx context.Context, host, index, shard, requestID string) (SimpleResponse, error) {
   111  	args := f.Called(ctx, host, index, shard, requestID)
   112  	return args.Get(0).(SimpleResponse), args.Error(1)
   113  }
   114  
   115  // Replica finder
   116  type fakeShardingState struct {
   117  	thisNode        string
   118  	ShardToReplicas map[string][]string
   119  	nodeResolver    *fakeNodeResolver
   120  }
   121  
   122  func newFakeShardingState(thisNode string, shardToReplicas map[string][]string, resolver *fakeNodeResolver) *fakeShardingState {
   123  	return &fakeShardingState{
   124  		thisNode:        thisNode,
   125  		ShardToReplicas: shardToReplicas,
   126  		nodeResolver:    resolver,
   127  	}
   128  }
   129  
   130  func (f *fakeShardingState) NodeName() string {
   131  	return f.thisNode
   132  }
   133  
   134  func (f *fakeShardingState) ResolveParentNodes(_ string, shard string) (map[string]string, error) {
   135  	replicas, ok := f.ShardToReplicas[shard]
   136  	if !ok {
   137  		return nil, fmt.Errorf("sharding state not found")
   138  	}
   139  
   140  	m := make(map[string]string)
   141  	for _, name := range replicas {
   142  		addr, _ := f.nodeResolver.NodeHostname(name)
   143  		m[name] = addr
   144  
   145  	}
   146  
   147  	return m, nil
   148  }
   149  
   150  // node resolver
   151  type fakeNodeResolver struct {
   152  	hosts map[string]string
   153  }
   154  
   155  func (r *fakeNodeResolver) NodeHostname(nodeName string) (string, bool) {
   156  	return r.hosts[nodeName], true
   157  }
   158  
   159  func newFakeNodeResolver(nodes []string) *fakeNodeResolver {
   160  	hosts := make(map[string]string)
   161  	for _, node := range nodes {
   162  		hosts[node] = node
   163  	}
   164  	return &fakeNodeResolver{hosts: hosts}
   165  }