Path: blob/main/SignalServiceKit/Messages/Attachments/V2/AttachmentStore/AttachmentStoreTests.swift
1 views
//
// Copyright 2024 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//
import Foundation
import XCTest
@testable import SignalServiceKit
class AttachmentStoreTests: XCTestCase {
private var db: InMemoryDB!
private var attachmentStore: AttachmentStore!
override func setUp() async throws {
db = InMemoryDB()
attachmentStore = AttachmentStore()
}
// MARK: - Inserts
func testInsert() throws {
var attachmentParams = Attachment.Record.mockPointer()
let referenceParams = AttachmentReference.ConstructionParams.mock(
owner: .thread(.globalThreadWallpaperImage(creationTimestamp: Date().ows_millisecondsSince1970)),
)
try db.write { tx in
_ = try attachmentStore.insert(
&attachmentParams,
reference: referenceParams,
tx: tx,
)
}
let (references, attachments) = db.read { tx in
let references = attachmentStore.fetchReferences(
owners: [.globalThreadWallpaperImage],
tx: tx,
)
let attachments = attachmentStore.fetch(
ids: references.map(\.attachmentRowId),
tx: tx,
)
return (references, attachments)
}
XCTAssertEqual(references.count, 1)
XCTAssertEqual(attachments.count, 1)
let reference = references[0]
let attachment = attachments[0]
XCTAssertEqual(reference.attachmentRowId, attachment.id)
assertEqual(attachmentParams, attachment)
assertEqual(referenceParams, reference)
}
func testMultipleInserts() throws {
let (threadId1, messageId1) = insertThreadAndInteraction()
let (threadId2, messageId2) = insertThreadAndInteraction()
let (threadId3, messageId3) = insertThreadAndInteraction()
let message1AttachmentIds: [UUID] = [.init()]
let message2AttachmentIds: [UUID] = [.init(), .init()]
let message3AttachmentIds: [UUID] = [.init(), .init(), .init()]
var attachmentIdToAttachmentParams = [UUID: Attachment.Record]()
var attachmentIdToAttachmentReferenceParams = [UUID: AttachmentReference.ConstructionParams]()
try db.write { tx in
for (messageId, threadId, attachmentIds) in [
(messageId1, threadId1, message1AttachmentIds),
(messageId2, threadId2, message2AttachmentIds),
(messageId3, threadId3, message3AttachmentIds),
] {
try attachmentIds.enumerated().forEach { index, id in
var attachmentParams = Attachment.Record.mockPointer()
let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId,
threadRowId: threadId,
orderInMessage: UInt32(index),
idInOwner: id,
)
try attachmentStore.insert(
&attachmentParams,
reference: attachmentReferenceParams,
tx: tx,
)
attachmentIdToAttachmentParams[id] = attachmentParams
attachmentIdToAttachmentReferenceParams[id] = attachmentReferenceParams
}
}
}
for (messageId, attachmentIds) in [
(messageId1, message1AttachmentIds),
(messageId2, message2AttachmentIds),
(messageId3, message3AttachmentIds),
] {
let (references, attachments) = db.read { tx in
let references = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId)],
tx: tx,
)
let attachments = attachmentStore.fetch(
ids: references.map(\.attachmentRowId),
tx: tx,
)
return (references, attachments)
}
XCTAssertEqual(references.count, attachmentIds.count)
XCTAssertEqual(attachments.count, attachmentIds.count)
for (reference, attachment) in zip(references, attachments) {
XCTAssertEqual(reference.attachmentRowId, attachment.id)
let attachmentId: UUID
switch reference.owner {
case .message(.bodyAttachment(let metadata)):
attachmentId = metadata.idInOwner!
default:
XCTFail("Unexpected owner type")
continue
}
guard
let attachmentParams = attachmentIdToAttachmentParams[attachmentId],
let referenceParams = attachmentIdToAttachmentReferenceParams[attachmentId]
else {
XCTFail("Unexpected attachment id")
continue
}
assertEqual(attachmentParams, attachment)
assertEqual(referenceParams, reference)
}
}
}
// MARK: -
func testInsertSameMediaName() {
let mediaName = Data(repeating: 27, count: 10).hexadecimalString
switch testAttachmentInsertError(
attachmentParams1: Attachment.Record.mockStream(streamInfo: .mock(mediaName: mediaName)),
attachmentParams2: Attachment.Record.mockStream(streamInfo: .mock(mediaName: mediaName)),
) {
case .duplicateMediaName:
break
case nil, .duplicatePlaintextHash:
XCTFail()
}
}
func testInsertSamePlaintextHash() throws {
let sha256ContentHash = UUID().data
switch testAttachmentInsertError(
attachmentParams1: Attachment.Record.mockStream(streamInfo: .mock(sha256ContentHash: sha256ContentHash)),
attachmentParams2: Attachment.Record.mockStream(streamInfo: .mock(sha256ContentHash: sha256ContentHash)),
) {
case .duplicatePlaintextHash:
break
case nil, .duplicateMediaName:
XCTFail()
}
}
private func testAttachmentInsertError(
attachmentParams1: Attachment.Record,
attachmentParams2: Attachment.Record,
) -> AttachmentInsertError? {
var attachmentParams1 = attachmentParams1
var attachmentParams2 = attachmentParams2
var attachmentInsertError: AttachmentInsertError?
let (threadId1, messageId1) = insertThreadAndInteraction()
let (threadId2, messageId2) = insertThreadAndInteraction()
db.write { tx in
let attachmentReferenceParams1 = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId1,
threadRowId: threadId1,
)
try! attachmentStore.insert(
&attachmentParams1,
reference: attachmentReferenceParams1,
tx: tx,
)
let message1References = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
)
let message1Attachment = attachmentStore.fetch(
ids: message1References.map(\.attachmentRowId),
tx: tx,
).first!
let attachmentReferenceParams2 = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId2,
threadRowId: threadId2,
)
do throws(AttachmentInsertError) {
try attachmentStore.insert(
&attachmentParams2,
reference: attachmentReferenceParams2,
tx: tx,
)
XCTFail("Should have thrown error!")
attachmentInsertError = nil
} catch {
attachmentInsertError = error
}
// Try again but insert using explicit owner adding.
attachmentStore.addReference(
attachmentReferenceParams2,
attachmentRowId: message1Attachment.id,
tx: tx,
)
}
let (
message1References,
message1Attachments,
message2References,
message2Attachments,
) = db.read { tx in
let message1References = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
)
let message1Attachments = attachmentStore.fetch(
ids: message1References.map(\.attachmentRowId),
tx: tx,
)
let message2References = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId2)],
tx: tx,
)
let message2Attachments = attachmentStore.fetch(
ids: message1References.map(\.attachmentRowId),
tx: tx,
)
return (message1References, message1Attachments, message2References, message2Attachments)
}
// Both messages should have one reference, to one attachment.
XCTAssertEqual(message1References.count, 1)
XCTAssertEqual(message1Attachments.count, 1)
XCTAssertEqual(message2References.count, 1)
XCTAssertEqual(message2Attachments.count, 1)
// But the attachments should be the same!
XCTAssertEqual(message1Attachments[0].id, message2Attachments[0].id)
// And it should have used the first attachment inserted.
XCTAssertEqual(message2Attachments[0].encryptionKey, attachmentParams1.encryptionKey)
return attachmentInsertError
}
// MARK: -
func testReinsertGlobalThreadAttachment() throws {
var attachmentParams1 = Attachment.Record.mockPointer()
let date1 = Date()
let referenceParams1 = AttachmentReference.ConstructionParams.mock(
owner: .thread(.globalThreadWallpaperImage(creationTimestamp: date1.ows_millisecondsSince1970)),
)
var attachmentParams2 = Attachment.Record.mockPointer()
let date2 = date1.addingTimeInterval(100)
let referenceParams2 = AttachmentReference.ConstructionParams.mock(
owner: .thread(.globalThreadWallpaperImage(creationTimestamp: date2.ows_millisecondsSince1970)),
)
try db.write { tx in
try attachmentStore.insert(
&attachmentParams1,
reference: referenceParams1,
tx: tx,
)
// Insert which should overwrite the existing row.
try attachmentStore.insert(
&attachmentParams2,
reference: referenceParams2,
tx: tx,
)
}
let (references, attachments) = db.read { tx in
let references = attachmentStore.fetchReferences(
owners: [.globalThreadWallpaperImage],
tx: tx,
)
let attachments = attachmentStore.fetch(
ids: references.map(\.attachmentRowId),
tx: tx,
)
return (references, attachments)
}
XCTAssertEqual(references.count, 1)
XCTAssertEqual(attachments.count, 1)
let reference = references[0]
let attachment = attachments[0]
XCTAssertEqual(reference.attachmentRowId, attachment.id)
assertEqual(attachmentParams2, attachment)
assertEqual(referenceParams2, reference)
}
// MARK: - Enumerate
func testEnumerateAttachmentReferences() throws {
let threadIdAndMessageIds = (0..<5).map { _ in
insertThreadAndInteraction()
}
// Insert many references to the same Params over and over.
var attachmentParams = Attachment.Record.mockStream()
let attachmentIdsInOwner: [UUID] = try db.write { tx in
var attachmentRowId: Attachment.IDType?
return try threadIdAndMessageIds.flatMap { threadId, messageId in
return try (0..<5).map { index in
let attachmentIdInOwner = UUID()
let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId,
threadRowId: threadId,
orderInMessage: UInt32(index),
idInOwner: attachmentIdInOwner,
)
if let attachmentRowId {
attachmentStore.addReference(
attachmentReferenceParams,
attachmentRowId: attachmentRowId,
tx: tx,
)
} else {
try attachmentStore.insert(
&attachmentParams,
reference: attachmentReferenceParams,
tx: tx,
)
attachmentRowId = tx.database.lastInsertedRowID
}
return attachmentIdInOwner
}
}
}
// Get the attachment id we just inserted.
let attachmentId = db.read { tx in
let references = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: threadIdAndMessageIds[0].interactionRowId)],
tx: tx,
)
let attachment = attachmentStore.fetch(
ids: references.map(\.attachmentRowId),
tx: tx,
)
return attachment.first!.id
}
// Insert some other references to other arbitrary attachments
try db.write { tx in
try threadIdAndMessageIds.forEach { threadId, messageId in
try (0..<5).forEach { index in
let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId,
threadRowId: threadId,
orderInMessage: UInt32(index),
)
var record = Attachment.Record.mockPointer()
try attachmentStore.insert(
&record,
reference: attachmentReferenceParams,
tx: tx,
)
}
}
}
// Check that we enumerate all the ids we created for the original attachment's id.
var enumeratedCount = 0
db.read { tx in
attachmentStore.enumerateAllReferences(
toAttachmentId: attachmentId,
tx: tx,
block: { reference, _ in
enumeratedCount += 1
XCTAssertEqual(reference.attachmentRowId, attachmentId)
switch reference.owner {
case .message(.bodyAttachment(let metadata)):
XCTAssertTrue(attachmentIdsInOwner.contains(metadata.idInOwner!))
default:
XCTFail("Unexpected attachment type!")
}
},
)
}
XCTAssertEqual(enumeratedCount, attachmentIdsInOwner.count)
}
func testOldestStickerPackReferences() throws {
// Sticker pack id to oldest row id.
var stickerPackIds = [Data: Int64]()
try db.write { tx in
let thread = insertThread(tx: tx)
let threadRowId = thread.sqliteRowId!
// Add some non sticker attachments.
for _ in 0..<10 {
let messageRowId = insertInteraction(thread: thread, tx: tx)
var record = Attachment.Record.mockStream()
try attachmentStore.insert(
&record,
reference: .mock(
owner: .message(.bodyAttachment(.init(
messageRowId: messageRowId,
receivedAtTimestamp: .random(in: 0..<9999999),
threadRowId: threadRowId,
contentType: .image,
isPastEditRevision: false,
caption: nil,
renderingFlag: .default,
orderInMessage: 0,
idInOwner: nil,
isViewOnce: false,
))),
),
tx: tx,
)
}
// Add some sticker attachments
for _ in 0..<10 {
let packId = UUID().data
// Each sticker pack gets multiple stickers, but we should
// dedupe down to the pack IDs.
let numStickersInPack = Int.random(in: 1...10)
for _ in 0..<numStickersInPack {
let stickerId = UInt32.random(in: 0..<UInt32.max)
let messageRowId = insertInteraction(thread: thread, tx: tx)
var record = Attachment.Record.mockStream()
try attachmentStore.insert(
&record,
reference: .mock(
owner: .message(.sticker(.init(
messageRowId: messageRowId,
receivedAtTimestamp: .random(in: 0..<9999999),
threadRowId: threadRowId,
contentType: .image,
isPastEditRevision: false,
stickerPackId: packId,
stickerId: stickerId,
))),
),
tx: tx,
)
stickerPackIds[packId] = min(messageRowId, stickerPackIds[packId] ?? .max)
}
}
}
let readPackReferences = db.read { tx in
attachmentStore.oldestStickerPackReferences(tx: tx)
}
XCTAssertEqual(readPackReferences.count, stickerPackIds.count)
readPackReferences.forEach { packRef in
XCTAssertEqual(packRef.messageRowId, stickerPackIds[packRef.stickerPackId])
}
}
func testallAttachmentIdsForSticker() throws {
// Sticker info to count of attachments with that info.
var stickerInfos = [StickerInfo: Int]()
try db.write { tx in
let thread = insertThread(tx: tx)
let threadRowId = thread.sqliteRowId!
// Add some non sticker attachments.
for _ in 0..<10 {
let messageRowId = insertInteraction(thread: thread, tx: tx)
var record = Attachment.Record.mockStream()
try attachmentStore.insert(
&record,
reference: .mock(
owner: .message(.bodyAttachment(.init(
messageRowId: messageRowId,
receivedAtTimestamp: .random(in: 0..<9999999),
threadRowId: threadRowId,
contentType: .image,
isPastEditRevision: false,
caption: nil,
renderingFlag: .default,
orderInMessage: 0,
idInOwner: nil,
isViewOnce: false,
))),
),
tx: tx,
)
}
// Add some sticker attachments
for _ in 0..<10 {
let packId = UUID().data
// Each sticker pack gets multiple stickers, but we should
// dedupe down to the pack IDs.
let numStickersInPack = Int.random(in: 1...10)
for _ in 0..<numStickersInPack {
let stickerId = UInt32.random(in: 0..<UInt32.max)
let attachmentCount = Int.random(in: 1...10)
stickerInfos[StickerInfo(
packId: packId,
packKey: Data(repeating: 8, count: Int(StickerManager.packKeyLength)),
stickerId: stickerId,
)] = attachmentCount
for _ in 0..<attachmentCount {
let messageRowId = insertInteraction(thread: thread, tx: tx)
var record = Attachment.Record.mockStream()
try attachmentStore.insert(
&record,
reference: .mock(
owner: .message(.sticker(.init(
messageRowId: messageRowId,
receivedAtTimestamp: .random(in: 0..<9999999),
threadRowId: threadRowId,
contentType: .image,
isPastEditRevision: false,
stickerPackId: packId,
stickerId: stickerId,
))),
),
tx: tx,
)
}
}
}
}
db.read { tx in
for (stickerInfo, expectedCount) in stickerInfos {
XCTAssertEqual(
expectedCount,
attachmentStore
.allAttachmentIdsForSticker(stickerInfo, tx: tx)
.count,
)
}
}
}
// MARK: - Update
func testMarkAsUploaded() throws {
let (threadId, messageId) = insertThreadAndInteraction()
try db.write { tx in
var attachmentParams = Attachment.Record.mockStream()
let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId,
threadRowId: threadId,
)
try attachmentStore.insert(
&attachmentParams,
reference: attachmentReferenceParams,
tx: tx,
)
}
func fetchAttachment() -> Attachment {
return db.read { tx in
let references = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId)],
tx: tx,
)
return attachmentStore.fetch(
ids: references.map(\.attachmentRowId),
tx: tx,
).first!
}
}
var attachment = fetchAttachment()
guard let stream = attachment.asStream() else {
XCTFail("Expected attachment stream!")
return
}
XCTAssertNil(attachment.latestTransitTierInfo)
let transitTierInfo = Attachment.TransitTierInfo(
cdnNumber: 3,
cdnKey: UUID().uuidString,
uploadTimestamp: Date().ows_millisecondsSince1970,
encryptionKey: UUID().data,
unencryptedByteCount: 100,
integrityCheck: .digestSHA256Ciphertext(UUID().data),
incrementalMacInfo: nil,
lastDownloadAttemptTimestamp: nil,
)
// Mark it as uploaded.
db.write { tx in
attachmentStore.saveLatestTransitTierInfo(
attachmentStream: stream,
transitTierInfo: transitTierInfo,
tx: tx,
)
}
// Refetch and check that it is appropriately marked.
attachment = fetchAttachment()
XCTAssertEqual(attachment.latestTransitTierInfo, transitTierInfo)
}
// MARK: - Remove Owner
func testRemoveOwner() throws {
let (threadId1, messageId1) = insertThreadAndInteraction()
let (threadId2, messageId2) = insertThreadAndInteraction()
// Create two references to the same attachment.
var attachmentParams = Attachment.Record.mockStream()
try db.write { tx in
try attachmentStore.insert(
&attachmentParams,
reference: AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId1,
threadRowId: threadId1,
),
tx: tx,
)
let attachmentId = tx.database.lastInsertedRowID
attachmentStore.addReference(
AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId2,
threadRowId: threadId2,
),
attachmentRowId: attachmentId,
tx: tx,
)
}
let (reference1, reference2) = db.read { tx in
let reference1 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
).first!
let reference2 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId2)],
tx: tx,
).first!
return (reference1, reference2)
}
XCTAssertEqual(reference1.attachmentRowId, reference2.attachmentRowId)
db.write { tx in
// Remove the first reference.
attachmentStore.removeReference(
reference: reference1,
tx: tx,
)
// The attachment should still exist.
XCTAssertNotNil(attachmentStore.fetch(
ids: [reference1.attachmentRowId],
tx: tx,
).first)
// Remove the second reference.
attachmentStore.removeReference(
reference: reference2,
tx: tx,
)
// The attachment should no longer exist.
XCTAssertNil(attachmentStore.fetch(
ids: [reference1.attachmentRowId],
tx: tx,
).first)
}
}
func testRemoveOwner_sameOwnerMultipleReferences() throws {
let (threadId1, messageId1) = insertThreadAndInteraction()
let referenceUUID1 = UUID()
let referenceUUID2 = UUID()
// Create two references to the same attachment on the same message.
var attachmentParams = Attachment.Record.mockStream()
try db.write { tx in
try attachmentStore.insert(
&attachmentParams,
reference: AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId1,
threadRowId: threadId1,
orderInMessage: 0,
idInOwner: referenceUUID1,
),
tx: tx,
)
let attachmentId = tx.database.lastInsertedRowID
attachmentStore.addReference(
AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageId1,
threadRowId: threadId1,
orderInMessage: 1,
idInOwner: referenceUUID2,
),
attachmentRowId: attachmentId,
tx: tx,
)
}
let (reference1, reference2) = db.read { tx in
let references = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
).sorted(by: {
$0.orderInOwningMessage! < $1.orderInOwningMessage!
})
return (references[0], references[1])
}
func checkIdInOwner(of reference: AttachmentReference, matches uuid: UUID) {
switch reference.owner {
case .message(let messageSource):
switch messageSource {
case .bodyAttachment(let metadata):
XCTAssertEqual(metadata.idInOwner, uuid)
default:
XCTFail("Unexpected reference type")
}
default:
XCTFail("Unexpected reference type")
}
}
checkIdInOwner(of: reference1, matches: referenceUUID1)
checkIdInOwner(of: reference2, matches: referenceUUID2)
db.write { tx in
// Remove the first reference.
attachmentStore.removeReference(
reference: reference1,
tx: tx,
)
// The attachment should still exist.
XCTAssertNotNil(attachmentStore.fetch(
ids: [reference1.attachmentRowId],
tx: tx,
).first)
// Refetch references
let references = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
)
XCTAssertEqual(references.count, 1)
checkIdInOwner(of: references[0], matches: referenceUUID2)
}
}
// MARK: Remove all thread owners
func testRemoveAllThreadOwners() throws {
var threadRowIds: [Int64?] = [nil]
db.write { tx in
for _ in 0..<5 {
threadRowIds.append(self.insertThread(tx: tx).sqliteRowId!)
}
}
try db.write { tx in
try threadRowIds.forEach { threadRowId in
var attachmentParams = Attachment.Record.mockPointer()
let timestamp = Date().ows_millisecondsSince1970
let referenceParams = AttachmentReference.ConstructionParams.mock(
owner: .thread(threadRowId.map {
.threadWallpaperImage(.init(threadRowId: $0, creationTimestamp: timestamp))
} ?? .globalThreadWallpaperImage(creationTimestamp: timestamp)),
)
try attachmentStore.insert(
&attachmentParams,
reference: referenceParams,
tx: tx,
)
}
}
func assertAttachmentCount(_ expectedCount: Int) {
db.read { tx in
let references = attachmentStore.fetchReferences(
owners: threadRowIds.map { threadRowId in
threadRowId.map {
.threadWallpaperImage(threadRowId: $0)
} ?? .globalThreadWallpaperImage
},
tx: tx,
)
let attachments = attachmentStore.fetch(
ids: references.map(\.attachmentRowId),
tx: tx,
)
XCTAssertEqual(references.count, expectedCount)
XCTAssertEqual(attachments.count, expectedCount)
}
}
assertAttachmentCount(threadRowIds.count)
// Remove all and count should be 0.
db.write { tx in
attachmentStore.removeAllThreadOwners(
tx: tx,
)
}
assertAttachmentCount(0)
}
// MARK: - Thread merging
func testThreadMerging() throws {
var threadId1: Int64!
var threadId2: Int64!
var threadId3: Int64!
var messageId1: Int64!
var messageId2: Int64!
var messageId3: Int64!
db.write { tx in
let thread1 = insertThread(tx: tx)
threadId1 = thread1.sqliteRowId!
messageId1 = insertInteraction(thread: thread1, tx: tx)
messageId2 = insertInteraction(thread: thread1, tx: tx)
let thread2 = insertThread(tx: tx)
threadId2 = thread2.sqliteRowId!
messageId3 = insertInteraction(thread: thread2, tx: tx)
let thread3 = insertThread(tx: tx)
threadId3 = thread3.sqliteRowId!
}
try db.write { tx in
for (threadRowId, messageRowId) in [
(threadId1, messageId1),
(threadId1, messageId2),
(threadId2, messageId3),
] {
// Create an attachment and reference for each message.
var attachmentParams = Attachment.Record.mockStream()
try attachmentStore.insert(
&attachmentParams,
reference: AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
messageRowId: messageRowId!,
threadRowId: threadRowId!,
),
tx: tx,
)
}
}
var (reference1, reference2, reference3) = db.read { tx in
let reference1 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
).first!
let reference2 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId2)],
tx: tx,
).first!
let reference3 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId3)],
tx: tx,
).first!
return (reference1, reference2, reference3)
}
func checkThreadRowId(of reference: AttachmentReference, matches threadId: Int64) {
switch reference.owner {
case .message(let messageSource):
switch messageSource {
case .bodyAttachment(let metadata):
XCTAssertEqual(metadata.threadRowId, threadId)
case .oversizeText(let metadata):
XCTAssertEqual(metadata.threadRowId, threadId)
case .linkPreview(let metadata):
XCTAssertEqual(metadata.threadRowId, threadId)
case .quotedReply(let metadata):
XCTAssertEqual(metadata.threadRowId, threadId)
case .sticker(let metadata):
XCTAssertEqual(metadata.threadRowId, threadId)
case .contactAvatar(let metadata):
XCTAssertEqual(metadata.threadRowId, threadId)
}
default:
XCTFail("Unexpected reference type")
}
}
checkThreadRowId(of: reference1, matches: threadId1)
checkThreadRowId(of: reference2, matches: threadId1)
checkThreadRowId(of: reference3, matches: threadId2)
db.write { tx in
// Merge threads
attachmentStore.updateMessageAttachmentThreadRowIdsForThreadMerge(
fromThreadRowId: threadId1,
intoThreadRowId: threadId3,
tx: tx,
)
}
(reference1, reference2, reference3) = db.read { tx in
let reference1 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId1)],
tx: tx,
).first!
let reference2 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId2)],
tx: tx,
).first!
let reference3 = attachmentStore.fetchReferences(
owners: [.messageBodyAttachment(messageRowId: messageId3)],
tx: tx,
).first!
return (reference1, reference2, reference3)
}
// The ones that were thread 1 are now thread 3.
checkThreadRowId(of: reference1, matches: threadId3)
checkThreadRowId(of: reference2, matches: threadId3)
checkThreadRowId(of: reference3, matches: threadId2)
}
// MARK: - UInt64 max values
/// Added as part of the introduction of `DBUInt64` and its first usage in
/// `Attachment.Record`. We don't need to implement this for every
/// subsequent type that uses `DBUInt64`.
func testUInt64MaxValues() throws {
try InMemoryDB().write { tx in
let attachment = Attachment.mock(
transitTierInfo: .mock(
uploadTimestamp: .max,
lastDownloadAttemptTimestamp: .max,
),
mediaTierInfo: .mock(
lastDownloadAttemptTimestamp: .max,
),
thumbnailInfo: .mock(
lastDownloadAttemptTimestamp: .max,
),
lastFullscreenViewTimestamp: .max,
)
var attachmentRecord = Attachment.Record(attachment: attachment)
try attachmentRecord.insert(tx.database)
let persisted = try Attachment.Record.fetchAll(tx.database)
XCTAssertEqual(persisted.count, 1)
XCTAssertEqual(persisted[0], attachmentRecord)
}
}
// MARK: - Helpers
private func insertThreadAndInteraction() -> (threadRowId: Int64, interactionRowId: Int64) {
return db.write { tx in
let thread = insertThread(tx: tx)
let interactionRowId = insertInteraction(thread: thread, tx: tx)
return (thread.sqliteRowId!, interactionRowId)
}
}
private func insertThread(tx: DBWriteTransaction) -> TSThread {
let thread = TSThread(uniqueId: UUID().uuidString)
try! thread.insert(tx.database)
return thread
}
private func insertInteraction(thread: TSThread, tx: DBWriteTransaction) -> Int64 {
let interaction = TSInteraction(timestamp: 0, receivedAtTimestamp: 0, thread: thread)
try! interaction.asRecord().insert(tx.database)
return interaction.sqliteRowId!
}
private func assertEqual(_ record: Attachment.Record, _ attachment: Attachment) {
var record = record
record.sqliteId = attachment.id
XCTAssertEqual(record, .init(attachment: attachment))
}
private func assertEqual(_ params: AttachmentReference.ConstructionParams, _ reference: AttachmentReference) {
switch (params.owner, reference.owner) {
case (.message, .message(let messageSource)):
let paramsRecord = AttachmentReference.MessageAttachmentReferenceRecord(
attachmentRowId: reference.attachmentRowId,
sourceFilename: params.sourceFilename,
sourceUnencryptedByteCount: params.sourceUnencryptedByteCount,
sourceMediaSizePixels: params.sourceMediaSizePixels,
messageSource: messageSource,
)
let referenceRecord = AttachmentReference.MessageAttachmentReferenceRecord(
attachmentReference: reference,
messageSource: messageSource,
)
XCTAssertEqual(paramsRecord, referenceRecord)
case (.storyMessage, .storyMessage(let storyMessageSource)):
let paramsRecord = AttachmentReference.StoryMessageAttachmentReferenceRecord(
attachmentRowId: reference.attachmentRowId,
sourceFilename: params.sourceFilename,
sourceUnencryptedByteCount: params.sourceUnencryptedByteCount,
sourceMediaSizePixels: params.sourceMediaSizePixels,
storyMessageSource: storyMessageSource,
)
let referenceRecord = AttachmentReference.StoryMessageAttachmentReferenceRecord(
attachmentReference: reference,
storyMessageSource: storyMessageSource,
)
XCTAssertEqual(paramsRecord, referenceRecord)
case (.thread, .thread(let threadSource)):
let paramsRecord = AttachmentReference.ThreadAttachmentReferenceRecord(
attachmentRowId: reference.attachmentRowId,
threadSource: threadSource,
)
let referenceRecord = AttachmentReference.ThreadAttachmentReferenceRecord(
attachmentRowId: reference.attachmentRowId,
threadSource: threadSource,
)
XCTAssertEqual(paramsRecord, referenceRecord)
case (.message, _), (.storyMessage, _), (.thread, _):
XCTFail("Non matching owner types")
}
}
}