Fix time-out issue in last "Index" stage

internal/stage_index_scan.go

@@ -4,7 +4,6 @@ import (
 	"database/sql"
 	"fmt"
 	"os"
-	"os/exec"
 	"path"
 	"sort"
 	"strings"
@@ -15,14 +14,60 @@ import (
 	"github.com/codecrafters-io/tester-utils/test_case_harness"
 )
 
-var testQueriesForCompanies = []string{
-	"SELECT id, name FROM companies WHERE country = 'micronesia'",
-	"SELECT id, name FROM companies WHERE country = 'north korea'",
-	"SELECT id, name FROM companies WHERE country = 'tonga'",
-	"SELECT id, name FROM companies WHERE country = 'eritrea'",
-	"SELECT id, name FROM companies WHERE country = 'republic of the congo'",
-	"SELECT id, name FROM companies WHERE country = 'montserrat'",
-	"SELECT id, name FROM companies WHERE country = 'chad'",
+var expectedQueryResultMap = map[string][]string{
+	"SELECT id, name FROM companies WHERE country = 'micronesia'": {
+		"1307865|college of micronesia",
+		"3696903|nanofabrica",
+		"4023193|fsm statistics",
+		"6132291|vital energy micronesia",
+		"6387751|fsm development bank",
+	},
+	"SELECT id, name FROM companies WHERE country = 'north korea'": {
+		"986681|isn network company limited",
+		"1573653|initial innovation limited",
+		"2828420|beacon point ltd",
+		"3485462|pyongyang university of science & technology (pust)",
+		"3969653|plastoform industries ltd",
+		"4271599|korea national insurance corporation",
+	},
+	"SELECT id, name FROM companies WHERE country = 'tonga'": {
+		"361142|tonga communications corporation",
+		"3186430|tonga development bank",
+		"3583436|leiola group limited",
+		"4796634|royco amalgamated company limited",
+		"7084593|tonga business enterprise centre",
+	},
+	"SELECT id, name FROM companies WHERE country = 'eritrea'": {
+		"121311|unilink s.c.",
+		"2102438|orange asmara it solutions",
+		"5729848|zara mining share company",
+		"6634629|asmara rental",
+	},
+	"SELECT id, name FROM companies WHERE country = 'republic of the congo'": {
+		"509721|skytic telecom",
+		"517263|somedia",
+		"2543747|its congo",
+		"2995059|petroleum trading congo e&p sa",
+	},
+	"SELECT id, name FROM companies WHERE country = 'montserrat'": {
+		"288999|government of montserrat",
+		"4472846|university of science, arts & technology",
+		"5316703|the abella group llc",
+	},
+	"SELECT id, name FROM companies WHERE country = 'chad'": {
+		"25661|ziyara",
+		"987266|hotel la mirande tchad",
+		"1313534|kreich avocats",
+		"2203192|societe des telecommunications du tchad",
+		"2435360|global logistics services limited (gls)",
+		"2466228|web tchad",
+		"2676248|hanana group",
+		"3775391|compagnie sucrière du tchad (cst)",
+		"4693857|wenaklabs",
+		"5021724|mariam high tech",
+		"5255614|bureau d'appui santé et environnement",
+		"6828605|tigo tchad",
+	},
 }
 
 func testIndexScan(stageHarness *test_case_harness.TestCaseHarness) error {
@@ -31,8 +76,8 @@ func testIndexScan(stageHarness *test_case_harness.TestCaseHarness) error {
 
 	_ = os.Remove("./test.db")
 
-	if err := exec.Command("cp", path.Join(os.Getenv("TESTER_DIR"), "companies.db"), "./test.db").Run(); err != nil {
-		logger.Errorf("Failed to create test database, this is a CodeCrafters error.")
+	if err := os.Link(path.Join(os.Getenv("TESTER_DIR"), "companies.db"), "./test.db"); err != nil {
+		logger.Errorf("Failed to create link for test database, this is a CodeCrafters error.")
 		return err
 	}
 
@@ -43,6 +88,10 @@ func testIndexScan(stageHarness *test_case_harness.TestCaseHarness) error {
 	}
 	defer db.Close()
 
+	var testQueriesForCompanies []string
+	for k := range expectedQueryResultMap {
+		testQueriesForCompanies = append(testQueriesForCompanies, k)
+	}
 	randomTestQueries := random.ShuffleArray(testQueriesForCompanies)[0:2]
 
 	for _, testQuery := range randomTestQueries {
@@ -58,12 +107,7 @@ func testIndexScan(stageHarness *test_case_harness.TestCaseHarness) error {
 
 		actualValues := splitBytesToLines(result.Stdout)
 
-		expectedValues, err := getExpectedValuesForQuery(db, testQuery)
-		if err != nil {
-			logger.Errorf("Failed to create test database, this is a CodeCrafters error.")
-			return err
-		}
-
+		expectedValues := expectedQueryResultMap[testQuery]
 		if len(actualValues) != len(expectedValues) {
 			return fmt.Errorf("Expected exactly %v lines of output, got: %v", len(expectedValues), len(actualValues))
 		}

internal/stage_index_scan_test.go

@@ -0,0 +1,69 @@
+package internal
+
+import (
+	"database/sql"
+	"fmt"
+	"sort"
+	"testing"
+
+	_ "modernc.org/sqlite"
+)
+
+func TestExpectedQueryResultMap(t *testing.T) {
+	db, err := sql.Open("sqlite", "../companies.db")
+	if err != nil {
+		t.Fatalf("Failed to open database: %v", err)
+	}
+	defer db.Close()
+
+	for query, expectedResults := range expectedQueryResultMap {
+		t.Run(fmt.Sprintf("query_%s", query), func(t *testing.T) {
+			// Execute the query
+			rows, err := db.Query(query)
+			if err != nil {
+				t.Fatalf("Failed to execute query '%s': %v", query, err)
+			}
+			defer rows.Close()
+
+			var actualResults []string
+			for rows.Next() {
+				var id int
+				var name string
+				if err := rows.Scan(&id, &name); err != nil {
+					t.Fatalf("Failed to scan row: %v", err)
+				}
+				actualResults = append(actualResults, fmt.Sprintf("%d|%s", id, name))
+			}
+
+			if err := rows.Err(); err != nil {
+				t.Fatalf("Error iterating rows: %v", err)
+			}
+
+			sort.Strings(expectedResults)
+			sort.Strings(actualResults)
+
+			if len(actualResults) != len(expectedResults) {
+				t.Errorf("Expected %d results, got %d results", len(expectedResults), len(actualResults))
+				t.Errorf("Expected: %v", expectedResults)
+				t.Errorf("Actual: %v", actualResults)
+				return
+			}
+
+			for i, expected := range expectedResults {
+				if i >= len(actualResults) {
+					t.Errorf("Missing result at index %d: expected %s", i, expected)
+					continue
+				}
+				if actualResults[i] != expected {
+					t.Errorf("Result mismatch at index %d: expected %s, got %s", i, expected, actualResults[i])
+				}
+			}
+
+			if t.Failed() {
+				t.Logf("Query: %s", query)
+				t.Logf("Expected results: %v", expectedResults)
+				t.Logf("Actual results: %v", actualResults)
+			}
+		})
+	}
+}