From d22576cbf24ea4a46d8b521360e6364e72b357c1 Mon Sep 17 00:00:00 2001 From: Tyson Gern Date: Thu, 2 May 2024 17:52:45 -0600 Subject: [PATCH] Extract CreateVector to testsupport --- internal/analyzer/analyze_test.go | 4 ++-- internal/analyzer/embeddings_gateway_test.go | 16 +++++----------- pkg/testsupport/vector_support.go | 7 +++++++ 3 files changed, 14 insertions(+), 13 deletions(-) create mode 100644 pkg/testsupport/vector_support.go diff --git a/internal/analyzer/analyze_test.go b/internal/analyzer/analyze_test.go index c464965..fbeda4e 100644 --- a/internal/analyzer/analyze_test.go +++ b/internal/analyzer/analyze_test.go @@ -13,7 +13,7 @@ import ( ) func TestAnalyzer_Analyze(t *testing.T) { - vector := createVector(0) + vector := testsupport.CreateVector(0) endpoint, server := testsupport.StartTestServer(t, func(mux *http.ServeMux) { testsupport.Handle(mux, "/embeddings", fmt.Sprintf(`{ "data": [ @@ -38,7 +38,7 @@ func TestAnalyzer_Analyze(t *testing.T) { err := a.Analyze(context.Background()) assert.NoError(t, err) - chunk1, err := embeddingsGateway.FindSimilar(createVector(0)) + chunk1, err := embeddingsGateway.FindSimilar(testsupport.CreateVector(0)) assert.NoError(t, err) assert.Equal(t, analyzer.CitedChunkRecord{Content: "chunk1", Source: "https://example.com"}, chunk1) } diff --git a/internal/analyzer/embeddings_gateway_test.go b/internal/analyzer/embeddings_gateway_test.go index cfc16f9..8de9397 100644 --- a/internal/analyzer/embeddings_gateway_test.go +++ b/internal/analyzer/embeddings_gateway_test.go @@ -20,7 +20,7 @@ func TestEmbeddingsGateway_UnprocessedIds(t *testing.T) { testDb.Execute("insert into data (id, source, content) values ('aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16', 'https://example.com', 'some content')") testDb.Execute("insert into chunks (id, data_id, content) values ('bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16', 'aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16','a chunk')") testDb.Execute("insert into chunks (id, data_id, content) values ('cccccccc-2f3f-4bc9-8dba-ba397156cc16', 'aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16','a chunk')") - vector := createVector(0) + vector := testsupport.CreateVector(0) testDb.Execute("insert into embeddings (chunk_id, embedding) values ('bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16', $1)", pgvector.NewVector(vector)) ids, err := gateway.UnprocessedIds() @@ -39,7 +39,7 @@ func TestEmbeddingsGateway_Save(t *testing.T) { testDb.Execute("insert into data (id, source, content) values ('aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16', 'https://example.com', 'some content')") testDb.Execute("insert into chunks (id, data_id, content) values ('bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16', 'aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16','a chunk')") - err := gateway.Save("bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16", createVector(0)) + err := gateway.Save("bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16", testsupport.CreateVector(0)) assert.NoError(t, err) chunkId, err := dbsupport.QueryOne(testDb.DB, "select chunk_id from embeddings", func(row *sql.Row, chunkId *string) error { @@ -60,19 +60,13 @@ func TestEmbeddingsGateway_FindSimilar(t *testing.T) { testDb.Execute("insert into chunks (id, data_id, content) values ('bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16', 'aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16','a chunk')") testDb.Execute("insert into chunks (id, data_id, content) values ('cccccccc-2f3f-4bc9-8dba-ba397156cc16', 'aaaaaaaa-2f3f-4bc9-8dba-ba397156cc16','another chunk')") - err := gateway.Save("bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16", createVector(0)) + err := gateway.Save("bbbbbbbb-2f3f-4bc9-8dba-ba397156cc16", testsupport.CreateVector(0)) assert.NoError(t, err) - err = gateway.Save("cccccccc-2f3f-4bc9-8dba-ba397156cc16", createVector(1)) + err = gateway.Save("cccccccc-2f3f-4bc9-8dba-ba397156cc16", testsupport.CreateVector(1)) assert.NoError(t, err) - similar, err := gateway.FindSimilar(createVector(1)) + similar, err := gateway.FindSimilar(testsupport.CreateVector(1)) assert.NoError(t, err) assert.Equal(t, analyzer.CitedChunkRecord{Content: "another chunk", Source: "https://example.com"}, similar) } - -func createVector(oneIndex int) []float32 { - embedding := make([]float32, 3072) - embedding[oneIndex] = 1 - return embedding -} diff --git a/pkg/testsupport/vector_support.go b/pkg/testsupport/vector_support.go new file mode 100644 index 0000000..1400f76 --- /dev/null +++ b/pkg/testsupport/vector_support.go @@ -0,0 +1,7 @@ +package testsupport + +func CreateVector(oneIndex int) []float32 { + embedding := make([]float32, 3072) + embedding[oneIndex] = 1 + return embedding +}