// Copyright (c) 2026 Lark Technologies Pte. Ltd. // SPDX-License-Identifier: MIT package draft import ( "strings" "testing" ) // buildSnapshotWithCard builds a minimal snapshot whose HTML body contains // a user section, a large attachment card, and optionally a quote block. func buildSnapshotWithCard(userContent, card, quote string) *DraftSnapshot { html := userContent + card + quote return &DraftSnapshot{ PrimaryHTMLPartID: "1", Body: &Part{ PartID: "1", MediaType: "text/html", Body: []byte(html), }, } } // buildSnapshotFromHTML wraps arbitrary HTML into a minimal snapshot. func buildSnapshotFromHTML(html string) *DraftSnapshot { return &DraftSnapshot{ PrimaryHTMLPartID: "1", Body: &Part{ PartID: "1", MediaType: "text/html", Body: []byte(html), }, } } const testLargeCard = `
original msg
old user content
`, testLargeCard, "") err := setBody(snap, `new user content
`, PatchOptions{}) if err != nil { t.Fatalf("setBody: %v", err) } newHTML := string(snap.Body.Body) if !strings.Contains(newHTML, "new user content") { t.Errorf("missing new content: %s", newHTML) } if strings.Contains(newHTML, "old user content") { t.Errorf("old content should be gone: %s", newHTML) } if !strings.Contains(newHTML, `id="large-file-area-123"`) { t.Errorf("card should be preserved: %s", newHTML) } if !strings.Contains(newHTML, "a.pdf") || !strings.Contains(newHTML, "tokA") { t.Errorf("card contents should be preserved: %s", newHTML) } } func TestSetBody_RespectsUserSuppliedCard(t *testing.T) { // When user's value already contains a large-file-area div, we must not // auto-duplicate. Result should have only the user's card, not the old one. snap := buildSnapshotWithCard(`old
`, testLargeCard, "") userCard := `new
`+userCard, PatchOptions{}) if err != nil { t.Fatalf("setBody: %v", err) } newHTML := string(snap.Body.Body) if !strings.Contains(newHTML, "userTok") { t.Errorf("user's card should be present: %s", newHTML) } if strings.Contains(newHTML, "large-file-area-123") { t.Errorf("old card should be gone (user supplied replacement): %s", newHTML) } // Should not be duplicated if strings.Count(newHTML, "large-file-area-") != 1 { t.Errorf("should have exactly one card, got %d: %s", strings.Count(newHTML, "large-file-area-"), newHTML) } } func TestSetBody_WithoutCardUnchangedBehavior(t *testing.T) { // No card in draft — setBody behaves as before. snap := &DraftSnapshot{ PrimaryHTMLPartID: "1", Body: &Part{ PartID: "1", MediaType: "text/html", Body: []byte(`old
`), }, } err := setBody(snap, `new
`, PatchOptions{}) if err != nil { t.Fatalf("setBody: %v", err) } if string(snap.Body.Body) != `new
` { t.Errorf("unexpected body: %q", string(snap.Body.Body)) } } func TestSetReplyBody_PreservesCardAndQuote(t *testing.T) { snap := buildSnapshotWithCard(`old user
`, testLargeCard, testQuoteBlock) err := setReplyBody(snap, `new user
`, PatchOptions{}) if err != nil { t.Fatalf("setReplyBody: %v", err) } newHTML := string(snap.Body.Body) if !strings.Contains(newHTML, "new user") { t.Errorf("missing new content: %s", newHTML) } if strings.Contains(newHTML, "old user") { t.Errorf("old user content should be gone: %s", newHTML) } if !strings.Contains(newHTML, `id="large-file-area-123"`) { t.Errorf("card should be preserved: %s", newHTML) } if !strings.Contains(newHTML, "original msg") { t.Errorf("quote should be preserved: %s", newHTML) } // Order: new user < card < quote newIdx := strings.Index(newHTML, "new user") cardIdx := strings.Index(newHTML, "large-file-area") quoteIdx := strings.Index(newHTML, "original msg") if !(newIdx < cardIdx && cardIdx < quoteIdx) { t.Errorf("expected order [user][card][quote]: newIdx=%d cardIdx=%d quoteIdx=%d, html=%s", newIdx, cardIdx, quoteIdx, newHTML) } } // TestSetReplyBody_ReplyToMessageWithCard verifies that when replying to // a message that itself contained a large attachment (so the quote block // in the draft contains the original sender's card), the user's own card // (sitting before the quote wrapper) is still preserved after // set_reply_body. The check in autoPreserveLargeAttachmentCard must only // look at value's user region, not inside the appended quote block. func TestSetReplyBody_ReplyToMessageWithCard(t *testing.T) { originalCardInQuote := `original message text
` + originalCardInQuote + `my old reply
`, testLargeCard, quoteWithOrigCard) err := setReplyBody(snap, `my new reply
`, PatchOptions{}) if err != nil { t.Fatalf("setReplyBody: %v", err) } newHTML := string(snap.Body.Body) // My card (from [my card] slot) should be preserved, even though the // quote block contains the original sender's card. if !strings.Contains(newHTML, `id="large-file-area-123"`) { t.Errorf("my own card (large-file-area-123) should be preserved: %s", newHTML) } // Original sender's card is still in the quote block (untouched by reply). if !strings.Contains(newHTML, `id="large-file-area-orig"`) { t.Errorf("original sender's card in quote should remain: %s", newHTML) } // New content present, old content gone. if !strings.Contains(newHTML, "my new reply") { t.Errorf("new content missing: %s", newHTML) } if strings.Contains(newHTML, "my old reply") { t.Errorf("old content should be gone: %s", newHTML) } // Order: new user content < my card < quote wrapper (which contains orig card) newIdx := strings.Index(newHTML, "my new reply") myCardIdx := strings.Index(newHTML, "large-file-area-123") quoteIdx := strings.Index(newHTML, "history-quote-wrapper") origCardIdx := strings.Index(newHTML, "large-file-area-orig") if !(newIdx < myCardIdx && myCardIdx < quoteIdx && quoteIdx < origCardIdx) { t.Errorf("expected order [user][my-card][quote[orig-card]]: new=%d my-card=%d quote=%d orig-card=%d\nhtml=%s", newIdx, myCardIdx, quoteIdx, origCardIdx, newHTML) } } func TestSetReplyBody_NoQuoteFallsBackToSetBody(t *testing.T) { // No quote — setReplyBody falls back to setBody, which preserves card. snap := buildSnapshotWithCard(`old
`, testLargeCard, "") err := setReplyBody(snap, `new
`, PatchOptions{}) if err != nil { t.Fatalf("setReplyBody: %v", err) } newHTML := string(snap.Body.Body) if !strings.Contains(newHTML, "large-file-area-123") { t.Errorf("card should be preserved: %s", newHTML) } if !strings.Contains(newHTML, "new") { t.Errorf("missing new content: %s", newHTML) } } func TestSplitAtLargeAttachment(t *testing.T) { cases := []struct { name string html string wantBefore string wantCardIn string // substring expected in card wantAfter string }{ { name: "no card", html: `hello
`, wantBefore: `hello
`, wantCardIn: "", wantAfter: "", }, { name: "card at end", html: `user
user
`, wantCardIn: "large-file-area-1", wantAfter: "", }, { name: "card before quote", html: `user
` + `user
`, wantCardIn: "large-file-area-1", wantAfter: `old user
` + testSigBlock) if err := setBody(snap, `new user
`, PatchOptions{}); err != nil { t.Fatalf("setBody: %v", err) } newHTML := string(snap.Body.Body) if !strings.Contains(newHTML, "new user") { t.Errorf("missing new content: %s", newHTML) } if !strings.Contains(newHTML, `class="lark-mail-signature"`) { t.Errorf("signature should be preserved: %s", newHTML) } if !strings.Contains(newHTML, "My Sig") { t.Errorf("signature content should be preserved: %s", newHTML) } // Order: new user content < signature newIdx := strings.Index(newHTML, "new user") sigIdx := strings.Index(newHTML, "lark-mail-signature") if newIdx > sigIdx { t.Errorf("signature should come after new content: new@%d sig@%d", newIdx, sigIdx) } } func TestSetBody_PreservesSignatureAndCard(t *testing.T) { snap := buildSnapshotFromHTML(`old
` + testSigBlock + testLargeCard) if err := setBody(snap, `new
`, PatchOptions{}); err != nil { t.Fatalf("setBody: %v", err) } newHTML := string(snap.Body.Body) newIdx := strings.Index(newHTML, "new") sigIdx := strings.Index(newHTML, "lark-mail-signature") cardIdx := strings.Index(newHTML, "large-file-area-123") if newIdx < 0 || sigIdx < 0 || cardIdx < 0 { t.Fatalf("missing parts: %s", newHTML) } if !(newIdx < sigIdx && sigIdx < cardIdx) { t.Errorf("expected order [new][sig][card], got new@%d sig@%d card@%d", newIdx, sigIdx, cardIdx) } } func TestSetBody_RespectsUserSuppliedSignature(t *testing.T) { snap := buildSnapshotFromHTML(`old
` + testSigBlock) userSig := `new
`+userSig, PatchOptions{}); err != nil { t.Fatalf("setBody: %v", err) } newHTML := string(snap.Body.Body) if !strings.Contains(newHTML, "User Sig") { t.Errorf("user-supplied sig should be present: %s", newHTML) } if strings.Contains(newHTML, "My Sig") { t.Errorf("old signature should be gone when user supplied their own: %s", newHTML) } // Only one signature wrapper if strings.Count(newHTML, "lark-mail-signature") != 1 { t.Errorf("expected exactly one signature wrapper, got %d", strings.Count(newHTML, "lark-mail-signature")) } } func TestSetReplyBody_PreservesSignatureAndQuote(t *testing.T) { snap := buildSnapshotFromHTML(`old user
` + testSigBlock + testQuoteBlock) if err := setReplyBody(snap, `new user
`, PatchOptions{}); err != nil { t.Fatalf("setReplyBody: %v", err) } newHTML := string(snap.Body.Body) newIdx := strings.Index(newHTML, "new user") sigIdx := strings.Index(newHTML, "lark-mail-signature") quoteIdx := strings.Index(newHTML, "history-quote-wrapper") if !(newIdx < sigIdx && sigIdx < quoteIdx) { t.Errorf("expected [new user][sig][quote], got new@%d sig@%d quote@%d", newIdx, sigIdx, quoteIdx) } } func TestSetReplyBody_PreservesAllThreeRegions(t *testing.T) { snap := buildSnapshotFromHTML(`old user
` + testSigBlock + testLargeCard + testQuoteBlock) if err := setReplyBody(snap, `new user
`, PatchOptions{}); err != nil { t.Fatalf("setReplyBody: %v", err) } newHTML := string(snap.Body.Body) newIdx := strings.Index(newHTML, "new user") sigIdx := strings.Index(newHTML, "lark-mail-signature") cardIdx := strings.Index(newHTML, "large-file-area-123") quoteIdx := strings.Index(newHTML, "history-quote-wrapper") if !(newIdx < sigIdx && sigIdx < cardIdx && cardIdx < quoteIdx) { t.Errorf("expected [new][sig][card][quote], got new@%d sig@%d card@%d quote@%d", newIdx, sigIdx, cardIdx, quoteIdx) } } // --------------------------------------------------------------------------- // ExtractSignatureBlock: symmetric with RemoveSignatureHTML // --------------------------------------------------------------------------- func TestExtractSignatureBlock_Symmetry(t *testing.T) { cases := []string{ `user
` + testSigBlock, `user
` + testSigBlock + testQuoteBlock, `user
` + testSigBlock + testLargeCard + testQuoteBlock, } for _, html := range cases { extracted := ExtractSignatureBlock(html) cleaned := RemoveSignatureHTML(html) if extracted == "" { t.Errorf("extract returned empty for: %s", html) continue } // The concatenation of cleaned + extracted (inserted back at the // right spot) should reconstitute the original. Since we don't // know the position, verify extract contains "lark-mail-signature" // and cleaned doesn't. if !strings.Contains(extracted, "lark-mail-signature") { t.Errorf("extract missing signature class: %s", extracted) } if strings.Contains(cleaned, "lark-mail-signature") { t.Errorf("clean still has signature: %s", cleaned) } // Length invariant: original == cleaned + extracted (bytes) if len(html) != len(cleaned)+len(extracted) { t.Errorf("length mismatch: %d != %d + %d", len(html), len(cleaned), len(extracted)) } } } func TestExtractSignatureBlock_NoSignature(t *testing.T) { if got := ExtractSignatureBlock(`just text
`); got != "" { t.Errorf("expected empty, got %q", got) } } func TestHTMLContainsLargeAttachment(t *testing.T) { cases := []struct { html string want bool }{ {`hello
`, false}, {``, true}, {`the text "large-file-area-" in body
`, false}, {`