larksuite--cli/shortcuts/mail/draft/large_attachment_parse_test.go

// Copyright (c) 2026 Lark Technologies Pte. Ltd.
// SPDX-License-Identifier: MIT

package draft

import (
	"encoding/base64"
	"encoding/json"
	"strings"
	"testing"
)

func TestParseLargeAttachmentTokens(t *testing.T) {
	encode := func(ids ...string) string {
		type item struct {
			ID string `json:"id"`
		}
		items := make([]item, len(ids))
		for i, id := range ids {
			items[i] = item{ID: id}
		}
		b, _ := json.Marshal(items)
		return base64.StdEncoding.EncodeToString(b)
	}

	cases := []struct {
		name    string
		headers []Header
		want    []string
	}{
		{
			name:    "empty headers",
			headers: nil,
			want:    nil,
		},
		{
			name:    "header present with one token",
			headers: []Header{{Name: LargeAttachmentIDsHeader, Value: encode("tokA")}},
			want:    []string{"tokA"},
		},
		{
			name:    "header present with multiple tokens in order",
			headers: []Header{{Name: LargeAttachmentIDsHeader, Value: encode("tokA", "tokB", "tokC")}},
			want:    []string{"tokA", "tokB", "tokC"},
		},
		{
			name:    "case-insensitive header name match",
			headers: []Header{{Name: "x-lms-large-attachment-ids", Value: encode("tokA")}},
			want:    []string{"tokA"},
		},
		{
			name:    "malformed base64 → nil",
			headers: []Header{{Name: LargeAttachmentIDsHeader, Value: "not!!base64"}},
			want:    nil,
		},
		{
			name:    "malformed JSON → nil",
			headers: []Header{{Name: LargeAttachmentIDsHeader, Value: base64.StdEncoding.EncodeToString([]byte("not json"))}},
			want:    nil,
		},
		{
			name:    "empty string IDs filtered out",
			headers: []Header{{Name: LargeAttachmentIDsHeader, Value: encode("tokA", "", "tokB")}},
			want:    []string{"tokA", "tokB"},
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got := parseLargeAttachmentTokens(tc.headers)
			if !equalStrings(got, tc.want) {
				t.Errorf("got %v, want %v", got, tc.want)
			}
		})
	}
}

func equalStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

func TestParseSizeDisplay(t *testing.T) {
	cases := []struct {
		in   string
		want int64
	}{
		{"25.0 MB", 26214400},
		{"1 GB", 1024 * 1024 * 1024},
		{"500 KB", 500 * 1024},
		{"42 B", 42},
		{"  25.0 MB  ", 26214400}, // whitespace tolerated
		{"25.0 mb", 26214400},     // case-insensitive
		{"garbage", 0},
		{"", 0},
		{"25", 0},    // no unit
		{"25 XB", 0}, // invalid unit
	}
	for _, tc := range cases {
		if got := parseSizeDisplay(tc.in); got != tc.want {
			t.Errorf("parseSizeDisplay(%q) = %d, want %d", tc.in, got, tc.want)
		}
	}
}

func TestParseLargeAttachmentItemsFromHTML(t *testing.T) {
	// Minimal HTML mirroring the structure generated by mail/large_attachment.go
	item := func(token, filename, size string) string {
		return `<div id="large-file-item">` +
			`<div><img src="x.png"/></div>` +
			`<div>` +
			`<div>` + filename + `</div>` +
			`<div><span>` + size + `</span></div>` +
			`</div>` +
			`<a href="x" data-mail-token="` + token + `">Download</a>` +
			`</div>`
	}

	html := `<div id="large-file-area-123">` +
		`<div>Title</div>` +
		item("tokA", "a.pdf", "25.0 MB") +
		item("tokB", "b.mov", "300 MB") +
		`</div>`

	got := ParseLargeAttachmentItemsFromHTML(html)
	if len(got) != 2 {
		t.Fatalf("expected 2 items, got %d: %+v", len(got), got)
	}
	if got["tokA"].FileName != "a.pdf" {
		t.Errorf("tokA filename: got %q, want %q", got["tokA"].FileName, "a.pdf")
	}
	if got["tokA"].SizeBytes != 26214400 {
		t.Errorf("tokA size: got %d, want 26214400", got["tokA"].SizeBytes)
	}
	if got["tokB"].FileName != "b.mov" {
		t.Errorf("tokB filename: got %q, want %q", got["tokB"].FileName, "b.mov")
	}
	if got["tokB"].SizeBytes != 300*1024*1024 {
		t.Errorf("tokB size: got %d, want %d", got["tokB"].SizeBytes, 300*1024*1024)
	}
}

func TestProjectLargeAttachments_MergesHeaderAndHTML(t *testing.T) {
	type idItem struct {
		ID string `json:"id"`
	}
	idsJSON, _ := json.Marshal([]idItem{{ID: "tokA"}, {ID: "tokB"}})
	headers := []Header{{Name: LargeAttachmentIDsHeader, Value: base64.StdEncoding.EncodeToString(idsJSON)}}

	html := `<div id="large-file-area-1">` +
		`<div>Title</div>` +
		`<div id="large-file-item"><div>a.pdf</div><div><span>25.0 MB</span></div><a data-mail-token="tokA">D</a></div>` +
		`<div id="large-file-item"><div>b.mov</div><div><span>300 MB</span></div><a data-mail-token="tokB">D</a></div>` +
		`</div>`

	got := projectLargeAttachments(headers, html)
	if len(got) != 2 {
		t.Fatalf("expected 2, got %d: %+v", len(got), got)
	}
	if got[0].Token != "tokA" || got[0].FileName != "a.pdf" || got[0].SizeBytes != 26214400 {
		t.Errorf("index 0: %+v", got[0])
	}
	if got[1].Token != "tokB" || got[1].FileName != "b.mov" {
		t.Errorf("index 1: %+v", got[1])
	}
}

func TestProjectLargeAttachments_HeaderWithoutHTML(t *testing.T) {
	// Token present in header but HTML missing the card entry (malformed draft).
	// We still report the token with empty meta.
	type idItem struct {
		ID string `json:"id"`
	}
	idsJSON, _ := json.Marshal([]idItem{{ID: "orphanToken"}})
	headers := []Header{{Name: LargeAttachmentIDsHeader, Value: base64.StdEncoding.EncodeToString(idsJSON)}}

	got := projectLargeAttachments(headers, "")
	if len(got) != 1 {
		t.Fatalf("expected 1, got %d", len(got))
	}
	if got[0].Token != "orphanToken" {
		t.Errorf("got token %q", got[0].Token)
	}
	if got[0].FileName != "" || got[0].SizeBytes != 0 {
		t.Errorf("expected empty meta, got %+v", got[0])
	}
}

func TestProjectLargeAttachments_NoHeader(t *testing.T) {
	got := projectLargeAttachments(nil, `<div id="large-file-area-1">...</div>`)
	if got != nil {
		t.Errorf("expected nil, got %v", got)
	}
}

func TestRemoveLargeAttachment_RemovesOneOfTwo(t *testing.T) {
	type idItem struct {
		ID string `json:"id"`
	}
	idsJSON, _ := json.Marshal([]idItem{{ID: "tokA"}, {ID: "tokB"}})
	headerValue := base64.StdEncoding.EncodeToString(idsJSON)
	html := `<html><body><p>hi</p><div id="large-file-area-1">` +
		`<div>Title</div>` +
		`<div id="large-file-item"><div>a.pdf</div><div><span>25.0 MB</span></div><a data-mail-token="tokA">D</a></div>` +
		`<div id="large-file-item"><div>b.mov</div><div><span>300 MB</span></div><a data-mail-token="tokB">D</a></div>` +
		`</div></body></html>`

	snapshot := &DraftSnapshot{
		Headers: []Header{{Name: LargeAttachmentIDsHeader, Value: headerValue}},
		Body: &Part{
			MediaType: "text/html",
			Body:      []byte(html),
		},
	}

	if err := removeLargeAttachment(snapshot, "tokA"); err != nil {
		t.Fatalf("removeLargeAttachment: %v", err)
	}

	// Header should contain only tokB
	tokens := parseLargeAttachmentTokens(snapshot.Headers)
	if !equalStrings(tokens, []string{"tokB"}) {
		t.Errorf("tokens after removal: got %v, want [tokB]", tokens)
	}

	// HTML should not contain data-mail-token="tokA" anymore, but still contain tokB and the container
	newHTML := string(snapshot.Body.Body)
	if strings.Contains(newHTML, `data-mail-token="tokA"`) {
		t.Errorf("HTML still contains tokA item:\n%s", newHTML)
	}
	if !strings.Contains(newHTML, `data-mail-token="tokB"`) {
		t.Errorf("HTML missing tokB item:\n%s", newHTML)
	}
	if !strings.Contains(newHTML, `id="large-file-area-1"`) {
		t.Errorf("HTML missing container (should still exist with tokB):\n%s", newHTML)
	}
}

func TestRemoveLargeAttachment_RemovesLastOneClearsContainer(t *testing.T) {
	type idItem struct {
		ID string `json:"id"`
	}
	idsJSON, _ := json.Marshal([]idItem{{ID: "tokOnly"}})
	headerValue := base64.StdEncoding.EncodeToString(idsJSON)
	html := `<html><body><p>hi</p><div id="large-file-area-1">` +
		`<div>Title</div>` +
		`<div id="large-file-item"><div>a.pdf</div><div><span>25.0 MB</span></div><a data-mail-token="tokOnly">D</a></div>` +
		`</div></body></html>`

	snapshot := &DraftSnapshot{
		Headers: []Header{{Name: LargeAttachmentIDsHeader, Value: headerValue}},
		Body: &Part{
			MediaType: "text/html",
			Body:      []byte(html),
		},
	}

	if err := removeLargeAttachment(snapshot, "tokOnly"); err != nil {
		t.Fatalf("removeLargeAttachment: %v", err)
	}

	// Header should be entirely removed (empty list)
	for _, h := range snapshot.Headers {
		if strings.EqualFold(h.Name, LargeAttachmentIDsHeader) {
			t.Errorf("header should have been removed when list is empty")
		}
	}

	// HTML should not contain the container at all
	newHTML := string(snapshot.Body.Body)
	if strings.Contains(newHTML, "large-file-area-1") {
		t.Errorf("container should have been removed:\n%s", newHTML)
	}
	if strings.Contains(newHTML, "large-file-item") {
		t.Errorf("no items should remain:\n%s", newHTML)
	}
	// Other body content should survive
	if !strings.Contains(newHTML, "<p>hi</p>") {
		t.Errorf("user body should remain:\n%s", newHTML)
	}
}

func TestRemoveLargeAttachment_UnknownToken(t *testing.T) {
	type idItem struct {
		ID string `json:"id"`
	}
	idsJSON, _ := json.Marshal([]idItem{{ID: "tokA"}})
	headerValue := base64.StdEncoding.EncodeToString(idsJSON)
	snapshot := &DraftSnapshot{
		Headers: []Header{{Name: LargeAttachmentIDsHeader, Value: headerValue}},
		Body:    &Part{MediaType: "text/html", Body: []byte(`<p>hi</p>`)},
	}
	err := removeLargeAttachment(snapshot, "unknown")
	if err == nil {
		t.Errorf("expected error for unknown token")
	}
}

func TestRemoveLargeAttachment_MissingHeader(t *testing.T) {
	snapshot := &DraftSnapshot{
		Headers: []Header{},
		Body:    &Part{MediaType: "text/html", Body: []byte(`<p>hi</p>`)},
	}
	err := removeLargeAttachment(snapshot, "any")
	if err == nil {
		t.Errorf("expected error when header is missing")
	}
}