fix(Add wg.Wait to GenerateColumnDescriptions)

Was missing the wg.Wait() for the API calls goroutine making the Groq API calls in GenerateColumnDescriptions
gwenwindflower · Apr 15, 2024 · c5c533a · c5c533a
1 parent d3d7f6e
commit c5c533a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/generate_column_desc.go b/generate_column_desc.go
@@ -136,6 +136,7 @@ func GenerateColumnDescriptions(tables shared.SourceTables) {
 			}(i, j)
 		}
 	}
+	wg.Wait()
 }
 
 func GetGroqResponse(prompt string) (GroqResponse, error) {

diff --git a/generate_column_desc_test.go b/generate_column_desc_test.go
@@ -9,13 +9,17 @@ import (
 func TestGetGroqResponse(t *testing.T) {
 	prompt := "Who destroyed Orthanc"
 	httpmock.Activate()
-	defer httpmock.Deactivate()
+	defer httpmock.DeactivateAndReset()
 	httpmock.RegisterResponder("POST", "https://api.groq.com/openai/v1/chat/completions",
 		httpmock.NewStringResponder(200, `{"choices": [{"index": 0, "message": {"role": "assistant","content": "Treebeard and the Ents destroyed Orthanc."}}]}`))
 	GroqResponse, err := GetGroqResponse(prompt)
 	if err != nil {
 		t.Error("expected", nil, "got", err)
 	}
+	info := httpmock.GetCallCountInfo()
+	if info["POST https://api.groq.com/openai/v1/chat/completions"] != 1 {
+		t.Error("expected", 1, "got", info["POST https://api.groq.com/openai/v1/chat/completions"])
+	}
 	expected := "Treebeard and the Ents destroyed Orthanc."
 	if GroqResponse.Choices[0].Message.Content != expected {
 		t.Error("expected", expected, "got", GroqResponse.Choices[0].Message.Content)
@@ -25,12 +29,17 @@ func TestGetGroqResponse(t *testing.T) {
 func TestGenerateColumnDescriptions(t *testing.T) {
 	ts := CreateTempSourceTables()
 	httpmock.Activate()
-	defer httpmock.Deactivate()
+	defer httpmock.DeactivateAndReset()
 	httpmock.RegisterResponder("POST", "https://api.groq.com/openai/v1/chat/completions",
 		httpmock.NewStringResponder(200, `{"choices": [{"index": 0, "message": {"role": "assistant","content": "lord of rivendell"}}]}`))
 	GenerateColumnDescriptions(ts)
-	expected := "lord of rivendell"
 
+	info := httpmock.GetCallCountInfo()
+	if info["POST https://api.groq.com/openai/v1/chat/completions"] != 2 {
+		t.Error("expected", 2, "got", info["POST https://api.groq.com/openai/v1/chat/completions"])
+	}
+
+	expected := "lord of rivendell"
 	desc := ts.SourceTables[0].Columns[0].Description
 	if desc != expected {
 		t.Error("expected", expected, "got", desc)
-Original file line number
+Diff line change
@@ Expand Up @@
     			}(i, j)
     		}
     	}
+    	wg.Wait()
     }
     func GetGroqResponse(prompt string) (GroqResponse, error) {
@@ Expand Down @@