github.com/weaviate/weaviate@v1.24.6/test/acceptance/replication/scale_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 replication
    13  
    14  import (
    15  	"context"
    16  	"fmt"
    17  	"testing"
    18  	"time"
    19  
    20  	"github.com/go-openapi/strfmt"
    21  	"github.com/stretchr/testify/assert"
    22  	"github.com/stretchr/testify/require"
    23  	"github.com/weaviate/weaviate/entities/models"
    24  	"github.com/weaviate/weaviate/entities/schema/crossref"
    25  	"github.com/weaviate/weaviate/test/docker"
    26  	"github.com/weaviate/weaviate/test/helper"
    27  	"github.com/weaviate/weaviate/test/helper/sample-schema/articles"
    28  	"github.com/weaviate/weaviate/usecases/replica"
    29  )
    30  
    31  func multiShardScaleOut(t *testing.T) {
    32  	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
    33  	defer cancel()
    34  
    35  	compose, err := docker.New().
    36  		WithWeaviateCluster().
    37  		WithText2VecContextionary().
    38  		Start(ctx)
    39  	require.Nil(t, err)
    40  	defer func() {
    41  		if err := compose.Terminate(ctx); err != nil {
    42  			t.Fatalf("failed to terminate test containers: %s", err.Error())
    43  		}
    44  	}()
    45  
    46  	helper.SetupClient(compose.GetWeaviate().URI())
    47  	paragraphClass := articles.ParagraphsClass()
    48  	paragraphClass.ShardingConfig = map[string]interface{}{
    49  		"desiredCount": 1,
    50  	}
    51  	articleClass := articles.ArticlesClass()
    52  	articleClass.ShardingConfig = map[string]interface{}{
    53  		"desiredCount": 1,
    54  	}
    55  
    56  	t.Run("create schema", func(t *testing.T) {
    57  		helper.CreateClass(t, paragraphClass)
    58  		helper.CreateClass(t, articleClass)
    59  	})
    60  
    61  	t.Run("insert paragraphs", func(t *testing.T) {
    62  		batch := make([]*models.Object, len(paragraphIDs))
    63  		for i, id := range paragraphIDs {
    64  			batch[i] = articles.NewParagraph().
    65  				WithID(id).
    66  				WithContents(fmt.Sprintf("paragraph#%d", i)).
    67  				Object()
    68  		}
    69  		createObjects(t, compose.GetWeaviate().URI(), batch)
    70  	})
    71  
    72  	t.Run("insert articles", func(t *testing.T) {
    73  		batch := make([]*models.Object, len(articleIDs))
    74  		for i, id := range articleIDs {
    75  			batch[i] = articles.NewArticle().
    76  				WithID(id).
    77  				WithTitle(fmt.Sprintf("Article#%d", i)).
    78  				Object()
    79  		}
    80  		createObjects(t, compose.GetWeaviateNode2().URI(), batch)
    81  	})
    82  
    83  	t.Run("add references", func(t *testing.T) {
    84  		refs := make([]*models.BatchReference, len(articleIDs))
    85  		for i := range articleIDs {
    86  			refs[i] = &models.BatchReference{
    87  				From: strfmt.URI(crossref.NewSource("Article", "hasParagraphs", articleIDs[i]).String()),
    88  				To:   strfmt.URI(crossref.NewLocalhost("Paragraph", paragraphIDs[i]).String()),
    89  			}
    90  		}
    91  		addReferences(t, compose.GetWeaviate().URI(), refs)
    92  	})
    93  
    94  	t.Run("scale out paragraphs", func(t *testing.T) {
    95  		c := getClass(t, compose.GetWeaviate().URI(), paragraphClass.Class)
    96  		c.ReplicationConfig.Factor = 2
    97  		updateClass(t, compose.GetWeaviate().URI(), c)
    98  	})
    99  
   100  	t.Run("assert paragraphs were scaled out", func(t *testing.T) {
   101  		n := getNodes(t, compose.GetWeaviate().URI())
   102  		var shardsFound int
   103  		for _, node := range n.Nodes {
   104  			for _, shard := range node.Shards {
   105  				if shard.Class == paragraphClass.Class {
   106  					assert.EqualValues(t, 10, shard.ObjectCount)
   107  					shardsFound++
   108  				}
   109  			}
   110  		}
   111  		assert.Equal(t, 2, shardsFound)
   112  	})
   113  
   114  	t.Run("scale out articles", func(t *testing.T) {
   115  		c := getClass(t, compose.GetWeaviate().URI(), articleClass.Class)
   116  		c.ReplicationConfig.Factor = 2
   117  		updateClass(t, compose.GetWeaviate().URI(), c)
   118  	})
   119  
   120  	t.Run("assert articles were scaled out", func(t *testing.T) {
   121  		n := getNodes(t, compose.GetWeaviate().URI())
   122  		var shardsFound int
   123  		for _, node := range n.Nodes {
   124  			for _, shard := range node.Shards {
   125  				if shard.Class == articleClass.Class {
   126  					assert.EqualValues(t, 10, shard.ObjectCount)
   127  					shardsFound++
   128  				}
   129  			}
   130  		}
   131  		assert.Equal(t, 2, shardsFound)
   132  	})
   133  
   134  	t.Run("kill a node and check contents of remaining node", func(t *testing.T) {
   135  		stopNode(ctx, t, compose, compose.GetWeaviateNode2().Name())
   136  		p := gqlGet(t, compose.GetWeaviate().URI(), paragraphClass.Class, replica.One)
   137  		assert.Len(t, p, 10)
   138  		a := gqlGet(t, compose.GetWeaviate().URI(), articleClass.Class, replica.One)
   139  		assert.Len(t, a, 10)
   140  	})
   141  }