github.com/weaviate/weaviate@v1.24.6/usecases/schema/restore_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 schema
    13  
    14  import (
    15  	"context"
    16  	"encoding/json"
    17  	"testing"
    18  
    19  	"github.com/stretchr/testify/assert"
    20  	"github.com/stretchr/testify/require"
    21  	"github.com/weaviate/weaviate/entities/backup"
    22  	"github.com/weaviate/weaviate/entities/models"
    23  	"github.com/weaviate/weaviate/usecases/sharding"
    24  )
    25  
    26  func TestRestoreClass_WithCircularRefs(t *testing.T) {
    27  	// When restoring a class, there could be circular refs between the classes,
    28  	// thus any validation that checks if linked classes exist would fail on the
    29  	// first class to import. Since we have no control over the order of imports
    30  	// when restoring, we need to relax this validation.
    31  
    32  	classes := []*models.Class{
    33  		{
    34  			Class: "Class_A",
    35  			Properties: []*models.Property{{
    36  				Name:     "to_Class_B",
    37  				DataType: []string{"Class_B"},
    38  			}, {
    39  				Name:     "to_Class_C",
    40  				DataType: []string{"Class_C"},
    41  			}},
    42  		},
    43  
    44  		{
    45  			Class: "Class_B",
    46  			Properties: []*models.Property{{
    47  				Name:     "to_Class_A",
    48  				DataType: []string{"Class_A"},
    49  			}, {
    50  				Name:     "to_Class_C",
    51  				DataType: []string{"Class_C"},
    52  			}},
    53  		},
    54  
    55  		{
    56  			Class: "Class_C",
    57  			Properties: []*models.Property{{
    58  				Name:     "to_Class_A",
    59  				DataType: []string{"Class_A"},
    60  			}, {
    61  				Name:     "to_Class_B",
    62  				DataType: []string{"Class_B"},
    63  			}},
    64  		},
    65  	}
    66  
    67  	mgr := newSchemaManager()
    68  
    69  	for _, classRaw := range classes {
    70  		schemaBytes, err := json.Marshal(classRaw)
    71  		require.Nil(t, err)
    72  
    73  		// for this particular test the sharding state does not matter, so we can
    74  		// initiate any new sharding state
    75  		shardingConfig, err := sharding.ParseConfig(nil, 1)
    76  		require.Nil(t, err)
    77  
    78  		nodes := fakeNodes{[]string{"node1", "node2"}}
    79  		shardingState, err := sharding.InitState(classRaw.Class, shardingConfig, nodes, 1, false)
    80  		require.Nil(t, err)
    81  
    82  		shardingBytes, err := shardingState.JSON()
    83  		require.Nil(t, err)
    84  
    85  		descriptor := backup.ClassDescriptor{Name: classRaw.Class, Schema: schemaBytes, ShardingState: shardingBytes}
    86  		err = mgr.RestoreClass(context.Background(), &descriptor, map[string]string{})
    87  		assert.Nil(t, err, "class passes validation")
    88  	}
    89  }
    90  
    91  func TestRestoreClass_WithNodeMapping(t *testing.T) {
    92  	classes := []*models.Class{{Class: "Class_A"}}
    93  
    94  	mgr := newSchemaManager()
    95  
    96  	for _, classRaw := range classes {
    97  		schemaBytes, err := json.Marshal(classRaw)
    98  		require.Nil(t, err)
    99  
   100  		shardingConfig, err := sharding.ParseConfig(nil, 2)
   101  		require.Nil(t, err)
   102  
   103  		nodes := fakeNodes{[]string{"node1", "node2"}}
   104  		shardingState, err := sharding.InitState(classRaw.Class, shardingConfig, nodes, 2, false)
   105  		require.Nil(t, err)
   106  
   107  		shardingBytes, err := shardingState.JSON()
   108  		require.Nil(t, err)
   109  
   110  		descriptor := backup.ClassDescriptor{Name: classRaw.Class, Schema: schemaBytes, ShardingState: shardingBytes}
   111  		err = mgr.RestoreClass(context.Background(), &descriptor, map[string]string{"node1": "new-node1"})
   112  		assert.NoError(t, err)
   113  
   114  		// Ensure that sharding state has been updated with the new node names
   115  		for _, shard := range mgr.ShardingState {
   116  			for _, v := range shard.Physical {
   117  				for _, node := range v.BelongsToNodes {
   118  					assert.Contains(t, []string{"new-node1", "node2"}, node)
   119  				}
   120  			}
   121  		}
   122  	}
   123  }
   124  
   125  type fakeNodes struct {
   126  	nodes []string
   127  }
   128  
   129  func (f fakeNodes) Candidates() []string {
   130  	return f.nodes
   131  }
   132  
   133  func (f fakeNodes) LocalName() string {
   134  	return f.nodes[0]
   135  }