package middleware import ( "net/http" "strconv" "strings" "testing" expect "github.com/yusing/go-proxy/internal/utils/testing" ) func TestInjectCSS(t *testing.T) { opts := OptionsRaw{ "target": "head", "html": "", } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html; charset=utf-8"}, }, respBody: []byte(` Test

Test

`), }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(` Test

Test

`)) contentLength, _ := strconv.Atoi(result.ResponseHeaders.Get("Content-Length")) expect.Equal(t, contentLength, len(result.Data), "Content-Length should be updated") } func TestInjectHTML_NonHTMLContent(t *testing.T) { opts := OptionsRaw{ "target": "head", "html": "", } originalBody := []byte(`{"message": "hello world"}`) result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"application/json"}, }, respBody: originalBody, }) expect.NoError(t, err) expect.Equal(t, result.Data, originalBody, "Non-HTML content should not be modified") } func TestInjectHTML_TargetNotFound(t *testing.T) { opts := OptionsRaw{ "target": ".nonexistent", "html": "
This should not appear
", } originalBody := []byte(` Test

Test

`) result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: originalBody, }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(string(originalBody)), "Content should remain unchanged when target not found") } func TestInjectHTML_MultipleTargets(t *testing.T) { opts := OptionsRaw{ "target": ".container", "html": "

Injected content

", } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(`
First container
Second container
`), }) expect.NoError(t, err) // Should only inject into the first matching element expectedContent := `
First container

Injected content

Second container
` expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(expectedContent)) } func TestInjectHTML_DifferentSelectors(t *testing.T) { testCases := []struct { name string selector string html string original string expected string }{ { name: "ID selector", selector: "#main", html: "By ID", original: `
Content
`, expected: `
ContentBy ID
`, }, { name: "Class selector", selector: ".highlight", html: "By class", original: `
Content
`, expected: `
ContentBy class
`, }, { name: "Element selector", selector: "body", html: "", original: `Content`, expected: `Content`, }, { name: "Attribute selector", selector: "[data-test='target']", html: "By attribute", original: `
Content
`, expected: `
ContentBy attribute
`, }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { opts := OptionsRaw{ "target": tc.selector, "html": tc.html, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(`` + tc.original + ``), }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(tc.expected)) }) } } func TestInjectHTML_EmptyInjection(t *testing.T) { opts := OptionsRaw{ "target": "head", "html": "", } originalBody := []byte(`Test`) result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: originalBody, }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(string(originalBody)), "Empty HTML injection should not change content") } func TestInjectHTML_ComplexHTML(t *testing.T) { opts := OptionsRaw{ "target": "body", "html": ``, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html; charset=utf-8"}, }, respBody: []byte(` Complex Page

Welcome

Some content here.

`), }) expect.NoError(t, err) resultStr := removeTabsAndNewlines(result.Data) expect.Equal(t, resultStr, removeTabsAndNewlines(` Complex Page

Welcome

Some content here.

`)) contentLength, _ := strconv.Atoi(result.ResponseHeaders.Get("Content-Length")) expect.Equal(t, contentLength, len(result.Data), "Content-Length should be updated correctly") } func TestInjectHTML_MalformedHTML(t *testing.T) { opts := OptionsRaw{ "target": "body", "html": "
Valid injection
", } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(`
Unclosed div

Some content`), }) expect.NoError(t, err) // Should handle malformed HTML gracefully expect.True(t, strings.Contains(string(result.Data), "Valid injection"), "Should inject content even with malformed HTML") } func TestInjectHTML_ContentTypes(t *testing.T) { testCases := []struct { name string contentType string shouldModify bool }{ {"HTML with charset", "text/html; charset=utf-8", true}, {"Plain HTML", "text/html", true}, {"XHTML", "application/xhtml+xml", true}, {"JSON", "application/json", false}, {"Plain text", "text/plain", false}, {"JavaScript", "application/javascript", false}, {"CSS", "text/css", false}, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { opts := OptionsRaw{ "target": "body", "html": "

Test injection
", } originalBody := []byte(`Original content`) result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{tc.contentType}, }, respBody: originalBody, }) expect.NoError(t, err) if tc.shouldModify { expect.True(t, strings.Contains(string(result.Data), "Test injection"), "Should modify HTML content for content-type: %s", tc.contentType) } else { expect.Equal(t, string(result.Data), string(originalBody), "Should not modify non-HTML content for content-type: %s", tc.contentType) } }) } } func TestInjectHTML_ReplaceTrue(t *testing.T) { opts := OptionsRaw{ "target": "body", "html": "
Replacement content
", "replace": true, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(` Test

Original content

More original content

`), }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(` Test
Replacement content
`)) contentLength, _ := strconv.Atoi(result.ResponseHeaders.Get("Content-Length")) expect.Equal(t, contentLength, len(result.Data), "Content-Length should be updated") } func TestInjectHTML_ReplaceVsAppend(t *testing.T) { originalBody := []byte(`
Original content
`) // Test append behavior (default) appendOpts := OptionsRaw{ "target": ".target", "html": "Added content", } appendResult, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: appendOpts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: originalBody, }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(appendResult.Data), removeTabsAndNewlines(`
Original contentAdded content
`)) // Test replace behavior replaceOpts := OptionsRaw{ "target": ".target", "html": "Replacement content", "replace": true, } replaceResult, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: replaceOpts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: originalBody, }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(replaceResult.Data), removeTabsAndNewlines(` Replacement content `)) } func TestInjectHTML_ReplaceWithDifferentSelectors(t *testing.T) { testCases := []struct { name string selector string original string html string expected string }{ { name: "ID selector replace", selector: "#main", original: `
Original content
`, html: "Replaced by ID", expected: `Replaced by ID`, }, { name: "Class selector replace", selector: ".highlight", original: `
Original content
`, html: "Replaced by class", expected: `Replaced by class`, }, { name: "Element selector replace", selector: "body", original: `Original content`, html: "
Replaced body content
", expected: `
Replaced body content
`, }, { name: "Attribute selector replace", selector: "[data-test='target']", original: `
Original content
`, html: "Replaced by attribute", expected: `Replaced by attribute`, }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { opts := OptionsRaw{ "target": tc.selector, "html": tc.html, "replace": true, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(`` + tc.original + ``), }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(tc.expected)) }) } } func TestInjectHTML_ReplaceWithEmpty(t *testing.T) { opts := OptionsRaw{ "target": ".content", "html": "", "replace": true, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(`
Content to be cleared
`), }) expect.NoError(t, err) expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(` `)) } func TestInjectHTML_ReplaceMultipleTargets(t *testing.T) { opts := OptionsRaw{ "target": ".container", "html": "

Replaced content

", "replace": true, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html"}, }, respBody: []byte(`
First container content
Second container content
`), }) expect.NoError(t, err) // Should only replace the first matching element expectedContent := `

Replaced content

Replaced content

` expect.Equal(t, removeTabsAndNewlines(result.Data), removeTabsAndNewlines(expectedContent)) } func TestInjectHTML_ReplaceComplexHTML(t *testing.T) { opts := OptionsRaw{ "target": "main", "html": `

New Section

This replaces the entire main content.

`, "replace": true, } result, err := newMiddlewareTest(ModifyHTML, &testArgs{ middlewareOpt: opts, respHeaders: http.Header{ "Content-Type": []string{"text/html; charset=utf-8"}, }, respBody: []byte(` Complex Page

Original Title

Original content that will be replaced.

More original content
`), }) expect.NoError(t, err) resultStr := removeTabsAndNewlines(result.Data) expect.Equal(t, resultStr, removeTabsAndNewlines(` Complex Page

New Section

This replaces the entire main content.

`)) contentLength, _ := strconv.Atoi(result.ResponseHeaders.Get("Content-Length")) expect.Equal(t, contentLength, len(result.Data), "Content-Length should be updated correctly") } func removeTabsAndNewlines[T string | []byte](s T) string { replacer := strings.NewReplacer( "\n", "", "\r", "", "\t", "", ) return replacer.Replace(string(s)) }