main_test.go

package main

import (
	"encoding/csv"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestMain(m *testing.M) {
	cleanupTestFiles()
	code := m.Run()
	cleanupTestFiles()
	os.Exit(code)
}

func cleanupTestFiles() {
	testFiles := []string{
		"test_data.csv",
		"test_data_part1.csv",
		"test_data_part2.csv",
		"test_data_part3.csv",
		"test_data_part4.csv",
		"test_data_part5.csv",
		"test_semicolon.csv",
		"test_semicolon_part1.csv",
		"test_semicolon_part2.csv",
		"test_tab.tsv",
		"test_tab_part1.tsv",
		"test_tab_part2.tsv",
		"test_tab_part3.tsv",
		"test_newlines.csv",
		"test_newlines_part1.csv",
		"test_newlines_part2.csv",
	}

	for _, file := range testFiles {
		os.Remove(file)
	}
}

func createTestCSV(filename string, rows [][]string, separator rune) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	writer := csv.NewWriter(file)
	writer.Comma = separator
	defer writer.Flush()

	for _, row := range rows {
		if err := writer.Write(row); err != nil {
			return err
		}
	}
	return nil
}

func readCSV(filename string, separator rune) ([][]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	reader := csv.NewReader(file)
	reader.Comma = separator
	return reader.ReadAll()
}

func TestBasicSplit(t *testing.T) {
	testData := [][]string{
		{"Name", "Age", "City"},
		{"John", "25", "New York"},
		{"Jane", "30", "Los Angeles"},
		{"Bob", "35", "Chicago"},
		{"Alice", "28", "Houston"},
		{"Charlie", "42", "Phoenix"},
		{"Diana", "33", "Philadelphia"},
	}

	err := createTestCSV("test_data.csv", testData, ',')
	if err != nil {
		t.Fatalf("Failed to create test CSV: %v", err)
	}

	os.Args = []string{"splitcsv", "-in", "test_data.csv"}

	part1, err := readCSV("test_data_part1.csv", ',')
	if err == nil {
		if len(part1) < 2 {
			t.Error("Part 1 should have header + data rows")
		}
		if part1[0][0] != "Name" {
			t.Error("Part 1 should have header row")
		}
	}
}

func TestSplitIntoMultipleParts(t *testing.T) {
	testData := [][]string{
		{"ID", "Value"},
	}
	for i := 1; i <= 10; i++ {
		testData = append(testData, []string{fmt.Sprintf("id%d", i), fmt.Sprintf("value%d", i)})
	}

	err := createTestCSV("test_data.csv", testData, ',')
	if err != nil {
		t.Fatalf("Failed to create test CSV: %v", err)
	}

	inputBase := strings.TrimSuffix(filepath.Base("test_data.csv"), filepath.Ext("test_data.csv"))
	inputExt := filepath.Ext("test_data.csv")
	inputDir := filepath.Dir("test_data.csv")

	expectedFilenames := []string{
		filepath.Join(inputDir, "test_data_part1.csv"),
		filepath.Join(inputDir, "test_data_part2.csv"),
		filepath.Join(inputDir, "test_data_part3.csv"),
	}

	for i := 0; i < 3; i++ {
		filename := fmt.Sprintf("%s_part%d%s", inputBase, i+1, inputExt)
		fullPath := filepath.Join(inputDir, filename)
		if fullPath != expectedFilenames[i] {
			t.Errorf("Expected filename %s, got %s", expectedFilenames[i], fullPath)
		}
	}
}

func TestRowDistribution(t *testing.T) {
	tests := []struct {
		totalRows int
		parts     int
		expected  []int
	}{
		{10, 2, []int{5, 5}},
		{10, 3, []int{4, 3, 3}},
		{11, 3, []int{4, 4, 3}},
		{7, 3, []int{3, 2, 2}},
		{8, 3, []int{3, 3, 2}},
	}

	for _, test := range tests {
		rowsPerPart := test.totalRows / test.parts
		extraRows := test.totalRows % test.parts

		distribution := make([]int, test.parts)
		for i := 0; i < test.parts; i++ {
			distribution[i] = rowsPerPart
			if extraRows > 0 {
				distribution[i]++
				extraRows--
			}
		}

		for i, expected := range test.expected {
			if distribution[i] != expected {
				t.Errorf("For %d rows in %d parts, part %d should have %d rows, got %d",
					test.totalRows, test.parts, i+1, expected, distribution[i])
			}
		}

		total := 0
		for _, count := range distribution {
			total += count
		}
		if total != test.totalRows {
			t.Errorf("Distribution doesn't add up: expected %d, got %d", test.totalRows, total)
		}
	}
}

func TestCustomSeparator(t *testing.T) {
	testData := [][]string{
		{"Name", "Age", "City"},
		{"John", "25", "New York"},
		{"Jane", "30", "Los Angeles"},
	}

	err := createTestCSV("test_semicolon.csv", testData, ';')
	if err != nil {
		t.Fatalf("Failed to create test CSV: %v", err)
	}

	data, err := readCSV("test_semicolon.csv", ';')
	if err != nil {
		t.Fatalf("Failed to read semicolon CSV: %v", err)
	}

	if len(data) != 3 {
		t.Errorf("Expected 3 rows, got %d", len(data))
	}

	if data[0][0] != "Name" {
		t.Errorf("Expected 'Name', got '%s'", data[0][0])
	}
}

func TestTabSeparator(t *testing.T) {
	testData := [][]string{
		{"Name", "Age", "City"},
		{"John", "25", "New York"},
		{"Jane", "30", "Los Angeles"},
	}

	err := createTestCSV("test_tab.tsv", testData, '\t')
	if err != nil {
		t.Fatalf("Failed to create test TSV: %v", err)
	}

	data, err := readCSV("test_tab.tsv", '\t')
	if err != nil {
		t.Fatalf("Failed to read tab TSV: %v", err)
	}

	if len(data) != 3 {
		t.Errorf("Expected 3 rows, got %d", len(data))
	}
}

func TestFilenameGeneration(t *testing.T) {
	tests := []struct {
		input    string
		parts    int
		expected []string
	}{
		{
			"data.csv",
			2,
			[]string{"data_part1.csv", "data_part2.csv"},
		},
		{
			"sales_2024.csv",
			3,
			[]string{"sales_2024_part1.csv", "sales_2024_part2.csv", "sales_2024_part3.csv"},
		},
		{
			"export.tsv",
			2,
			[]string{"export_part1.tsv", "export_part2.tsv"},
		},
		{
			"/path/to/file.csv",
			2,
			[]string{"/path/to/file_part1.csv", "/path/to/file_part2.csv"},
		},
	}

	for _, test := range tests {
		inputBase := strings.TrimSuffix(filepath.Base(test.input), filepath.Ext(test.input))
		inputExt := filepath.Ext(test.input)
		inputDir := filepath.Dir(test.input)

		for i := 0; i < test.parts; i++ {
			filename := fmt.Sprintf("%s_part%d%s", inputBase, i+1, inputExt)
			fullPath := filepath.Join(inputDir, filename)

			if fullPath != test.expected[i] {
				t.Errorf("For input %s part %d, expected %s, got %s",
					test.input, i+1, test.expected[i], fullPath)
			}
		}
	}
}

func TestQuotedFields(t *testing.T) {
	// Test CSV with quoted fields containing separators
	testData := [][]string{
		{"Name", "Description", "Price"},
		{"John", "Software Engineer, Senior", "75,000"},
		{"Jane", "Manager with \"experience\"", "85,500"},
		{"Bob", "Sales Rep, handles big accounts", "45,000"},
	}

	err := createTestCSV("test_quoted.csv", testData, ',')
	if err != nil {
		t.Fatalf("Failed to create quoted CSV: %v", err)
	}
	defer os.Remove("test_quoted.csv")

	data, err := readCSV("test_quoted.csv", ',')
	if err != nil {
		t.Fatalf("Failed to read quoted CSV: %v", err)
	}

	if len(data) != 4 {
		t.Errorf("Expected 4 rows, got %d", len(data))
	}

	// Check that quoted field with comma is preserved as single field
	if data[1][1] != "Software Engineer, Senior" {
		t.Errorf("Expected 'Software Engineer, Senior', got '%s'", data[1][1])
	}

	// Check that quoted field with quotes is preserved
	if data[2][1] != "Manager with \"experience\"" {
		t.Errorf("Expected 'Manager with \"experience\"', got '%s'", data[2][1])
	}

	// Check that numeric field with comma is preserved
	if data[1][2] != "75,000" {
		t.Errorf("Expected '75,000', got '%s'", data[1][2])
	}
}

func TestIsEOF(t *testing.T) {
	tests := []struct {
		err      error
		expected bool
	}{
		{fmt.Errorf("EOF"), true},
		{fmt.Errorf("other error"), false},
		{fmt.Errorf("unexpected EOF"), false},
	}

	for _, test := range tests {
		result := isEOF(test.err)
		if result != test.expected {
			t.Errorf("For error '%v', expected %v, got %v", test.err, test.expected, result)
		}
	}
}

func TestDetectQuotedFields(t *testing.T) {
	tests := []struct {
		input    string
		sep      rune
		expected []bool
	}{
		{
			`"quoted","unquoted","another quoted"`,
			',',
			[]bool{true, true, true},
		},
		{
			`unquoted,"quoted field",unquoted`,
			',',
			[]bool{false, true, false},
		},
		{
			`"field with ""escaped"" quotes","normal field"`,
			',',
			[]bool{true, true},
		},
		{
			`"quoted with, comma",unquoted,"another quoted"`,
			',',
			[]bool{true, false, true},
		},
		{
			`field1;field2;"quoted field"`,
			';',
			[]bool{false, false, true},
		},
		{
			`"",,"non-empty"`,
			',',
			[]bool{true, false, true},
		},
		{
			`simple,fields,only`,
			',',
			[]bool{false, false, false},
		},
		{
			`"all","fields","quoted"`,
			',',
			[]bool{true, true, true},
		},
	}

	for i, test := range tests {
		result := detectQuotedFields(test.input, test.sep)
		if len(result) != len(test.expected) {
			t.Errorf("Test %d: expected %d fields, got %d", i+1, len(test.expected), len(result))
			continue
		}

		for j, expected := range test.expected {
			if j < len(result) && result[j] != expected {
				t.Errorf("Test %d, field %d: expected %v, got %v", i+1, j+1, expected, result[j])
			}
		}
	}
}

func TestCleanNewlinesInQuoted(t *testing.T) {
	tests := []struct {
		record       []string
		quotedFields []bool
		expected     []string
	}{
		{
			[]string{"field1", "field with\nnewline", "field3"},
			[]bool{false, true, false},
			[]string{"field1", "field with newline", "field3"},
		},
		{
			[]string{"no newlines", "also no newlines"},
			[]bool{false, true},
			[]string{"no newlines", "also no newlines"},
		},
		{
			[]string{"field\nwith\nnewlines", "quoted\r\nfield", "normal"},
			[]bool{false, true, false},
			[]string{"field\nwith\nnewlines", "quoted field", "normal"},
		},
		{
			[]string{"multiple\n\nlines", "quoted\r\n\nfield"},
			[]bool{true, true},
			[]string{"multiple  lines", "quoted  field"},
		},
		{
			[]string{"field1", "field2", "field3"},
			[]bool{false, false, false},
			[]string{"field1", "field2", "field3"},
		},
	}

	for i, test := range tests {
		result := cleanNewlinesInQuoted(test.record, test.quotedFields)
		if len(result) != len(test.expected) {
			t.Errorf("Test %d: expected %d fields, got %d", i+1, len(test.expected), len(result))
			continue
		}

		for j, expected := range test.expected {
			if result[j] != expected {
				t.Errorf("Test %d, field %d: expected '%s', got '%s'", i+1, j+1, expected, result[j])
			}
		}
	}
}

func createTestCSVWithNewlines(filename string, content string) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.WriteString(content)
	return err
}

func TestNewlineHandlingIntegration(t *testing.T) {
	// Create CSV with newlines in quoted fields
	csvContent := `name,description,price
"Product A","This is a long
description with line breaks",100.50
Product B,Simple description,200
"Product C","Another item with
multiple
line breaks",150.75
Regular Product,No newlines here,75`

	err := createTestCSVWithNewlines("test_newlines.csv", csvContent)
	if err != nil {
		t.Fatalf("Failed to create test CSV with newlines: %v", err)
	}

	// Test our functions directly first
	rawLine1 := `"Product A","This is a long
description with line breaks",100.50`
	quotedFields1 := detectQuotedFields(rawLine1, ',')
	expectedQuoted1 := []bool{true, true, false}

	if len(quotedFields1) != len(expectedQuoted1) {
		t.Errorf("Expected %d quoted field flags, got %d", len(expectedQuoted1), len(quotedFields1))
	} else {
		for i, expected := range expectedQuoted1 {
			if quotedFields1[i] != expected {
				t.Errorf("Field %d: expected quoted=%v, got %v", i+1, expected, quotedFields1[i])
			}
		}
	}

	// Test that we can read the created file back
	data, err := readCSV("test_newlines.csv", ',')
	if err != nil {
		t.Fatalf("Failed to read newlines CSV: %v", err)
	}

	if len(data) != 5 {
		t.Errorf("Expected 5 rows (including header), got %d", len(data))
	}

	// Check that header is correct
	if data[0][0] != "name" || data[0][1] != "description" || data[0][2] != "price" {
		t.Errorf("Header row incorrect: got %v", data[0])
	}

	// Check first product with newlines - the CSV reader should preserve newlines in quoted fields
	if data[1][0] != "Product A" {
		t.Errorf("Expected 'Product A', got '%s'", data[1][0])
	}

	// The description should contain a newline since it was in quotes
	expectedDesc := "This is a long\ndescription with line breaks"
	if data[1][1] != expectedDesc {
		t.Errorf("Expected '%s', got '%s'", expectedDesc, data[1][1])
	}

	// Test our cleaning function on this record
	cleanedRecord := cleanNewlinesInQuoted(data[1], []bool{true, true, false})
	expectedCleanedDesc := "This is a long description with line breaks"
	if cleanedRecord[1] != expectedCleanedDesc {
		t.Errorf("Cleaned description should be '%s', got '%s'", expectedCleanedDesc, cleanedRecord[1])
	}

	// The unquoted field should remain unchanged
	if cleanedRecord[2] != data[1][2] {
		t.Errorf("Unquoted price field should remain unchanged: expected '%s', got '%s'", data[1][2], cleanedRecord[2])
	}
}

func BenchmarkRowDistribution(b *testing.B) {
	for i := 0; i < b.N; i++ {
		totalRows := 1000000
		parts := 7

		rowsPerPart := totalRows / parts
		extraRows := totalRows % parts

		for j := 0; j < parts; j++ {
			count := rowsPerPart
			if extraRows > 0 {
				count++
				extraRows--
			}
			_ = count
		}
	}
}