grokkit/cmd/docs_test.go

package cmd

import (
	"strings"
	"testing"
)

func TestBuildDocsMessages(t *testing.T) {
	tests := []struct {
		language   string
		code       string
		styleCheck string
	}{
		{language: "Go", code: "package main\nfunc Foo() {}", styleCheck: "godoc"},
		{language: "Python", code: "def foo():\n    pass", styleCheck: "PEP 257"},
		{language: "C", code: "int foo(void) { return 0; }", styleCheck: "Doxygen"},
		{language: "C++", code: "int foo() { return 0; }", styleCheck: "Doxygen"},
		{language: "JavaScript", code: "function foo() {}", styleCheck: "JSDoc"},
		{language: "TypeScript", code: "function foo(): void {}", styleCheck: "JSDoc"},
		{language: "Rust", code: "pub fn foo() {}", styleCheck: "rustdoc"},
		{language: "Ruby", code: "def foo; end", styleCheck: "YARD"},
		{language: "Java", code: "public void foo() {}", styleCheck: "Javadoc"},
		{language: "Shell", code: "foo() { echo hello; }", styleCheck: "shell comment"},
	}

	for _, tt := range tests {
		t.Run(tt.language, func(t *testing.T) {
			msgs := buildDocsMessages(tt.language, tt.code)

			if len(msgs) != 2 {
				t.Fatalf("expected 2 messages, got %d", len(msgs))
			}
			if msgs[0]["role"] != "system" {
				t.Errorf("first message role = %q, want %q", msgs[0]["role"], "system")
			}
			if msgs[1]["role"] != "user" {
				t.Errorf("second message role = %q, want %q", msgs[1]["role"], "user")
			}
			if !strings.Contains(msgs[0]["content"], tt.styleCheck) {
				t.Errorf("system prompt missing %q for language %s; got: %s",
					tt.styleCheck, tt.language, msgs[0]["content"])
			}
			if !strings.Contains(msgs[1]["content"], tt.code) {
				t.Errorf("user prompt missing code for language %s", tt.language)
			}
		})
	}
}

func TestDocStyle(t *testing.T) {
	tests := []struct {
		language string
		want     string
	}{
		{"go", "godoc"},
		{"Go", "godoc"},
		{"python", "PEP 257 docstring"},
		{"c", "Doxygen"},
		{"c++", "Doxygen"},
		{"javascript", "JSDoc"},
		{"typescript", "JSDoc"},
		{"rust", "rustdoc"},
		{"ruby", "YARD"},
		{"java", "Javadoc"},
		{"shell", "shell comment"},
		{"bash", "shell comment"},
		{"unknown", "standard documentation comment"},
	}

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