dealspace/lib/watermark_test.go

package lib

import (
	"archive/zip"
	"bytes"
	"image"
	"image/color"
	"image/png"
	"testing"
)

// TestWatermarkPDF tests PDF watermarking with a minimal PDF.
func TestWatermarkPDF(t *testing.T) {
	// Minimal valid PDF (empty page)
	minimalPDF := []byte(`%PDF-1.4
1 0 obj
<< /Type /Catalog /Pages 2 0 R >>
endobj
2 0 obj
<< /Type /Pages /Kids [3 0 R] /Count 1 >>
endobj
3 0 obj
<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] >>
endobj
xref
0 4
0000000000 65535 f 
0000000009 00000 n 
0000000058 00000 n 
0000000115 00000 n 
trailer
<< /Size 4 /Root 1 0 R >>
startxref
192
%%EOF`)

	label := "CONFIDENTIAL - Test User - 2026-02-28 - TestProject"

	out, err := WatermarkPDF(minimalPDF, label)
	if err != nil {
		t.Fatalf("WatermarkPDF failed: %v", err)
	}

	// Basic checks
	if len(out) == 0 {
		t.Fatal("WatermarkPDF returned empty output")
	}

	// Watermarked PDF should be larger than original
	if len(out) <= len(minimalPDF) {
		t.Errorf("Watermarked PDF (%d bytes) should be larger than original (%d bytes)", len(out), len(minimalPDF))
	}

	// Should still start with PDF header
	if !bytes.HasPrefix(out, []byte("%PDF")) {
		t.Error("Output doesn't look like a PDF (missing %PDF header)")
	}
}

// TestWatermarkImagePNG tests image watermarking with PNG.
func TestWatermarkImagePNG(t *testing.T) {
	// Create a minimal test PNG (100x100 red square)
	img := image.NewRGBA(image.Rect(0, 0, 100, 100))
	red := color.RGBA{255, 0, 0, 255}
	for y := 0; y < 100; y++ {
		for x := 0; x < 100; x++ {
			img.Set(x, y, red)
		}
	}

	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatalf("Failed to create test PNG: %v", err)
	}
	originalPNG := buf.Bytes()

	label := "CONFIDENTIAL - Test User - 2026-02-28"

	out, err := WatermarkImage(originalPNG, "image/png", label)
	if err != nil {
		t.Fatalf("WatermarkImage failed: %v", err)
	}

	// Basic checks
	if len(out) == 0 {
		t.Fatal("WatermarkImage returned empty output")
	}

	// Verify it's still a valid PNG
	_, err = png.Decode(bytes.NewReader(out))
	if err != nil {
		t.Errorf("Output is not a valid PNG: %v", err)
	}
}

// TestWatermarkDOCX tests DOCX watermarking with a minimal document.
func TestWatermarkDOCX(t *testing.T) {
	// Create a minimal DOCX (which is a ZIP file with specific structure)
	minimalDOCX := createMinimalDOCX()

	label := "CONFIDENTIAL - Test User - 2026-02-28 - TestProject"

	out, err := WatermarkDOCX(minimalDOCX, label)
	if err != nil {
		t.Fatalf("WatermarkDOCX failed: %v", err)
	}

	// Basic checks
	if len(out) == 0 {
		t.Fatal("WatermarkDOCX returned empty output")
	}

	// Watermarked DOCX should be larger (we added header)
	if len(out) <= len(minimalDOCX) {
		t.Errorf("Watermarked DOCX (%d bytes) should be larger than original (%d bytes)", len(out), len(minimalDOCX))
	}

	// Should still be a valid ZIP
	if !bytes.HasPrefix(out, []byte("PK")) {
		t.Error("Output doesn't look like a DOCX/ZIP (missing PK header)")
	}
}

// TestWatermarkDispatch tests the main dispatch function.
func TestWatermarkDispatch(t *testing.T) {
	tests := []struct {
		name     string
		mimeType string
		data     []byte
	}{
		{"unknown type passthrough", "application/octet-stream", []byte("hello world")},
		{"text passthrough", "text/plain", []byte("hello world")},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			out, err := Watermark(tt.data, tt.mimeType, "test label")
			if err != nil {
				t.Errorf("Watermark failed: %v", err)
			}
			// Unknown types should pass through unchanged
			if !bytes.Equal(out, tt.data) {
				t.Error("Unknown MIME type should pass through unchanged")
			}
		})
	}
}

// TestWatermarkEmptyInput tests error handling for empty inputs.
func TestWatermarkEmptyInput(t *testing.T) {
	_, err := WatermarkPDF(nil, "test")
	if err == nil {
		t.Error("WatermarkPDF should fail on nil input")
	}

	_, err = WatermarkImage(nil, "image/png", "test")
	if err == nil {
		t.Error("WatermarkImage should fail on nil input")
	}

	_, err = WatermarkDOCX(nil, "test")
	if err == nil {
		t.Error("WatermarkDOCX should fail on nil input")
	}
}

// createMinimalDOCX creates a minimal valid DOCX for testing.
func createMinimalDOCX() []byte {
	var buf bytes.Buffer
	w := zip.NewWriter(&buf)

	// [Content_Types].xml
	contentTypes := `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">
  <Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>
  <Default Extension="xml" ContentType="application/xml"/>
  <Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/>
</Types>`
	f, _ := w.Create("[Content_Types].xml")
	f.Write([]byte(contentTypes))

	// _rels/.rels
	rels := `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
  <Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/>
</Relationships>`
	f, _ = w.Create("_rels/.rels")
	f.Write([]byte(rels))

	// word/document.xml
	doc := `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">
  <w:body>
    <w:p>
      <w:r>
        <w:t>Test Document</w:t>
      </w:r>
    </w:p>
  </w:body>
</w:document>`
	f, _ = w.Create("word/document.xml")
	f.Write([]byte(doc))

	// word/_rels/document.xml.rels
	docRels := `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">
</Relationships>`
	f, _ = w.Create("word/_rels/document.xml.rels")
	f.Write([]byte(docRels))

	w.Close()
	return buf.Bytes()
}