grokkit/cmd/testgen_test.go

package cmd

import (
	"strings"
	"testing"
)

func TestRemoveSourceComments(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name  string
		input string
		want  string
		lang  string
	}{
		{
			name: "no comments",
			input: `package cmd
import "testing"

func Foo() {}`,
			want: `package cmd
import "testing"

func Foo() {}`,
			lang: "Go",
		},
		{
			name: "last modified",
			input: `// Last modified: 2026-03-02
package cmd`,
			want: `package cmd`,
			lang: "Go",
		},
		{
			name: "generated by",
			input: `// Generated by grokkit testgen
package cmd`,
			want: `package cmd`,
			lang: "Go",
		},
		{
			name: "multiple removable lines",
			input: `line1
// Last modified: foo
line3
// Generated by: bar
line5`,
			want: `line1
line3
line5`,
			lang: "Go",
		},
		{
			name: "partial match no remove",
			input: `// Modified something else
package cmd`,
			want: `// Modified something else
package cmd`,
			lang: "Go",
		},
		{
			name: "python testgen",
			input: `# testgen: generated
def foo(): pass`,
			want: `def foo(): pass`,
			lang: "Python",
		},
		{
			name: "c testgen",
			input: `/* testgen */
int foo() {}`,
			want: `int foo() {}`,
			lang: "C",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := removeSourceComments(tt.input, tt.lang)
			if got != tt.want {
				t.Errorf("removeSourceComments() =\n%q\nwant\n%q", got, tt.want)
			}
		})
	}
}

func TestGetTestPrompt(t *testing.T) {
	t.Parallel()

	tests := []struct {
		lang       string
		wantPrefix string
	}{
		{"Go", "You are an expert Go testing specialist."},
		{"Python", "You are a pytest expert."},
		{"C", "You are a C unit testing expert using Check framework."},
		{"C++", "You are a Google Test expert."},
		{"Invalid", ""},
	}

	for _, tt := range tests {
		t.Run(tt.lang, func(t *testing.T) {
			got := getTestPrompt(tt.lang)
			if tt.wantPrefix != "" && !strings.HasPrefix(got, tt.wantPrefix) {
				t.Errorf("getTestPrompt(%q) prefix =\n%q\nwant %q", tt.lang, got[:100], tt.wantPrefix)
			}
			if tt.wantPrefix == "" && got != "" {
				t.Errorf("getTestPrompt(%q) = %q, want empty", tt.lang, got)
			}
		})
	}
}

func TestGetTestFilePath(t *testing.T) {
	t.Parallel()

	tests := []struct {
		filePath string
		lang     string
		want     string
	}{
		{"foo.go", "Go", "foo_test.go"},
		{"dir/foo.py", "Python", "dir/test_foo.py"},
		{"bar.c", "C", "test_bar.c"},
		{"baz.cpp", "C++", "test_baz.cpp"},
	}

	for _, tt := range tests {
		t.Run(tt.filePath+"_"+tt.lang, func(t *testing.T) {
			got := getTestFilePath(tt.filePath, tt.lang)
			if got != tt.want {
				t.Errorf("getTestFilePath(%q, %q) = %q, want %q", tt.filePath, tt.lang, got, tt.want)
			}
		})
	}
}

func TestGetCodeLang(t *testing.T) {
	t.Parallel()

	tests := []struct {
		lang string
		want string
	}{
		{"Go", "go"},
		{"Python", "python"},
		{"C", "c"},
		{"C++", "c"},
	}

	for _, tt := range tests {
		t.Run(tt.lang, func(t *testing.T) {
			got := getCodeLang(tt.lang)
			if got != tt.want {
				t.Errorf("getCodeLang(%q) = %q, want %q", tt.lang, got, tt.want)
			}
		})
	}
}