llmgw/internal/web/codeblocks_test.go

package web

import (
	"strings"
	"testing"
)

func TestSafePathRefusesEverythingThatEscapes(t *testing.T) {
	// The name comes from a language model and goes into an archive somebody
	// extracts. These are things to expect, not things to hope against.
	for _, bad := range []string{
		"../secrets.txt",
		"../../.ssh/authorized_keys",
		"a/../../b",
		"..",
		".",
		"",
		"   ",
		"a/./../../b",
		"foo\x00.go",
		"foo\n../bar",
		strings.Repeat("a", 200) + ".go",
		strings.Repeat("a/", 20) + "x.go",
	} {
		if got := safePath(bad); got != "" {
			t.Errorf("safePath(%q) = %q; it must be refused", bad, got)
		}
	}
}

func TestSafePathKeepsOrdinaryNames(t *testing.T) {
	for in, want := range map[string]string{
		"main.go":               "main.go",
		"cmd/server/main.go":    "cmd/server/main.go",
		"./pkg/x.go":            "pkg/x.go",
		"/src/app.py":           "src/app.py",
		`src\app.py`:            "src/app.py",
		"a/b/../c.txt":          "a/c.txt",
		"Dockerfile":            "Dockerfile",
		"deploy/k8s/deploy.yml": "deploy/k8s/deploy.yml",
	} {
		if got := safePath(in); got != want {
			t.Errorf("safePath(%q) = %q, want %q", in, got, want)
		}
	}
}

func TestExtractNamesFilesTheWayModelsWriteThem(t *testing.T) {
	reply := "Here you go.\n\n" +
		"```go:cmd/api/main.go\npackage main\n```\n\n" +
		"**internal/store/store.go**\n" +
		"```go\npackage store\n```\n\n" +
		"```python scripts/run.py\nprint(1)\n```\n\n" +
		"```Dockerfile\nFROM scratch\n```\n"

	files := ExtractCode(reply)
	if len(files) != 4 {
		t.Fatalf("got %d files: %+v", len(files), files)
	}
	want := []string{"cmd/api/main.go", "internal/store/store.go", "scripts/run.py", "Dockerfile"}
	for i, w := range want {
		if files[i].Path != w {
			t.Errorf("file %d: path %q, want %q", i, files[i].Path, w)
		}
	}
	if !strings.Contains(files[0].Body, "package main") {
		t.Errorf("body lost: %q", files[0].Body)
	}
}

func TestAnUnnamedBlockIsStillReturned(t *testing.T) {
	// Losing someone's code because the model did not label it is worse than
	// naming it badly.
	files := ExtractCode("```go\npackage main\n```\n```\nplain text\n```\n")
	if len(files) != 2 {
		t.Fatalf("got %d files", len(files))
	}
	if files[0].Path != "snippet-1.go" || files[1].Path != "snippet-2.txt" {
		t.Errorf("generated names wrong: %q, %q", files[0].Path, files[1].Path)
	}
}

func TestAnEscapingNameFallsBackRatherThanBeingRepaired(t *testing.T) {
	// Silently rewriting ../../x.go into x.go would put a file somewhere the
	// reader did not expect. It gets a generated name instead.
	files := ExtractCode("```go:../../evil.go\npackage main\n```\n")
	if len(files) != 1 {
		t.Fatalf("got %d files", len(files))
	}
	if strings.Contains(files[0].Path, "..") {
		t.Fatalf("an escaping path survived: %q", files[0].Path)
	}
	if files[0].Path != "snippet-1.go" {
		t.Errorf("path = %q, want the generated name", files[0].Path)
	}
}

func TestTwoBlocksClaimingOnePathBothSurvive(t *testing.T) {
	// Models show a file, then show it again with one line changed.
	files := ExtractCode("```go:main.go\nv1\n```\n```go:main.go\nv2\n```\n")
	if len(files) != 2 {
		t.Fatalf("got %d files", len(files))
	}
	if files[0].Path == files[1].Path {
		t.Fatalf("one overwrote the other: both are %q", files[0].Path)
	}
	if files[1].Path != "main-2.go" {
		t.Errorf("second path = %q, want main-2.go", files[1].Path)
	}
}

func TestProseMentioningAFileIsNotAFilename(t *testing.T) {
	reply := "You should now edit main.go and rebuild.\n```go\npackage main\n```\n"
	files := ExtractCode(reply)
	if len(files) != 1 {
		t.Fatalf("got %d files", len(files))
	}
	if files[0].Path != "snippet-1.go" {
		t.Errorf("a sentence was read as a filename: %q", files[0].Path)
	}
}

func TestATextReplyWithNoCodeYieldsNothing(t *testing.T) {
	if files := ExtractCode("No code here, just words about `main.go`.\n"); len(files) != 0 {
		t.Errorf("got %d files from prose: %+v", len(files), files)
	}
}

func TestTheNumberOfFilesIsBounded(t *testing.T) {
	var b strings.Builder
	for range maxCodeFiles + 20 {
		b.WriteString("```go\nx\n```\n")
	}
	if got := len(ExtractCode(b.String())); got > maxCodeFiles {
		t.Errorf("returned %d files, over the %d cap", got, maxCodeFiles)
	}
}

func TestAPathCommentInsideTheBlockNamesTheFile(t *testing.T) {
	// Exactly what qwen2.5-coder-3b returned when asked for a two-file Go
	// server: a plain ```go fence with the path as the first line of the file.
	// Both files landed as snippet-1.go and snippet-2.go before this.
	reply := "```go\n// cmd/server/main.go\n\npackage main\n\nfunc main() {}\n```\n\n" +
		"```go\n// internal/handler/handler.go\n\npackage handler\n```\n"
	files := ExtractCode(reply)
	if len(files) != 2 {
		t.Fatalf("got %d files", len(files))
	}
	if files[0].Path != "cmd/server/main.go" || files[1].Path != "internal/handler/handler.go" {
		t.Fatalf("paths lost: %q, %q", files[0].Path, files[1].Path)
	}
	// The comment stays: people write those anyway, and rewriting model output
	// to tidy a filename is a surprise nobody asked for.
	if !strings.Contains(files[0].Body, "// cmd/server/main.go") {
		t.Error("the path comment was stripped from the body")
	}
}

func TestPathCommentsAcrossLanguages(t *testing.T) {
	for _, tc := range []struct{ block, want string }{
		{"```python\n# scripts/run.py\nprint(1)\n```\n", "scripts/run.py"},
		{"```sql\n-- db/schema.sql\nSELECT 1;\n```\n", "db/schema.sql"},
		{"```html\n<!-- web/index.html -->\n<p>hi</p>\n```\n", "web/index.html"},
		{"```c\n/* src/main.c */\nint main(){}\n```\n", "src/main.c"},
		{"```go\n\n// pkg/x.go\npackage x\n```\n", "pkg/x.go"},
	} {
		files := ExtractCode(tc.block)
		if len(files) != 1 {
			t.Fatalf("%q: got %d files", tc.block, len(files))
		}
		if files[0].Path != tc.want {
			t.Errorf("%q: path %q, want %q", tc.block, files[0].Path, tc.want)
		}
	}
}

func TestAnOrdinaryCommentIsNotAPath(t *testing.T) {
	// A sentence that happens to mention a file is not a filename, and a
	// comment below the first couple of lines is just a comment.
	for _, block := range []string{
		"```go\n// this reads config.json at startup\npackage main\n```\n",
		"```go\npackage main\n\n// internal/x.go\nfunc main(){}\n```\n",
		"```go\n// Package main implements the server.\npackage main\n```\n",
	} {
		files := ExtractCode(block)
		if len(files) != 1 {
			t.Fatalf("got %d files", len(files))
		}
		if files[0].Path != "snippet-1.go" {
			t.Errorf("%q was read as a filename in:\n%s", files[0].Path, block)
		}
	}
}