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)
	}
}