// 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 `
` } html := `hi
` 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 := `hi
hi
") { 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(`hi
`)}, } 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(`hi
`)}, } err := removeLargeAttachment(snapshot, "any") if err == nil { t.Errorf("expected error when header is missing") } }