document_test.go

Documentation: gitlab.hexacode.org/go-libs/chromem-go

     1  package chromem
     2  
     3  import (
     4  	"context"
     5  	"reflect"
     6  	"testing"
     7  )
     8  
     9  func TestDocument_New(t *testing.T) {
    10  	ctx := context.Background()
    11  	id := "test"
    12  	metadata := map[string]string{"foo": "bar"}
    13  	vectors := []float32{-0.40824828, 0.40824828, 0.81649655} // normalized version of `{-0.1, 0.1, 0.2}`
    14  	content := "hello world"
    15  	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
    16  		return vectors, nil
    17  	}
    18  
    19  	tt := []struct {
    20  		name          string
    21  		id            string
    22  		metadata      map[string]string
    23  		vectors       []float32
    24  		content       string
    25  		embeddingFunc EmbeddingFunc
    26  	}{
    27  		{
    28  			name:          "No embedding",
    29  			id:            id,
    30  			metadata:      metadata,
    31  			vectors:       nil,
    32  			content:       content,
    33  			embeddingFunc: embeddingFunc,
    34  		},
    35  		{
    36  			name:          "With embedding",
    37  			id:            id,
    38  			metadata:      metadata,
    39  			vectors:       vectors,
    40  			content:       content,
    41  			embeddingFunc: embeddingFunc,
    42  		},
    43  	}
    44  
    45  	for _, tc := range tt {
    46  		t.Run(tc.name, func(t *testing.T) {
    47  			// Create document
    48  			d, err := NewDocument(ctx, id, metadata, vectors, content, embeddingFunc)
    49  			if err != nil {
    50  				t.Fatal("expected no error, got", err)
    51  			}
    52  			// We can compare with DeepEqual after removing the embedding function
    53  			d.Embedding = nil
    54  			exp := Document{
    55  				ID:       id,
    56  				Metadata: metadata,
    57  				Content:  content,
    58  			}
    59  			if !reflect.DeepEqual(exp, d) {
    60  				t.Fatalf("expected %+v, got %+v", exp, d)
    61  			}
    62  		})
    63  	}
    64  }
    65
View as plain text