Path: blob/main/SignalServiceKit/Backups/Attachments/BackupListMediaManagerTests.swift
1 views
//
// Copyright 2025 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//
import GRDB
import LibSignalClient
import Testing
@testable import SignalServiceKit
public class BackupListMediaManagerTests {
private lazy var accountKeyStore = AccountKeyStore(backupSettingsStore: backupSettingsStore)
private let attachmentStore = AttachmentStore()
private let attachmentUploadStore = AttachmentUploadStore()
private let backupAttachmentDownloadProgress = BackupAttachmentDownloadProgressMock()
private let backupAttachmentDownloadStore = BackupAttachmentDownloadStore()
private let backupAttachmentUploadProgress = BackupAttachmentUploadProgressMock(initialCompleted: 0, total: 100)
private let backupAttachmentUploadStore = BackupAttachmentUploadStore()
private let backupAttachmentUploadEraStore = BackupAttachmentUploadEraStore()
private let backupListMediaStore = BackupListMediaStore()
private lazy var backupMediaErrorNotificationPresenter = BackupMediaErrorNotificationPresenter(
dateProvider: dateProvider,
db: db,
notificationPresenter: notificationPresenter,
)
private let backupRequestManager = BackupRequestManagerMock()
private let backupSettingsStore = BackupSettingsStore()
private let dateProvider: DateProvider = { Date() }
private let db = InMemoryDB()
private let interactionStore = InteractionStoreImpl()
private let notificationPresenter = NoopNotificationPresenterImpl()
private let orphanedBackupAttachmentStore = OrphanedBackupAttachmentStore()
private let remoteConfigManager = StubbableRemoteConfigManager()
private let tsAccountManager = MockTSAccountManager()
private lazy var backupAttachmentUploadScheduler = BackupAttachmentUploadScheduler(
attachmentStore: attachmentStore,
backupAttachmentUploadStore: backupAttachmentUploadStore,
backupAttachmentUploadEraStore: backupAttachmentUploadEraStore,
dateProvider: dateProvider,
interactionStore: interactionStore,
remoteConfigProvider: remoteConfigManager,
tsAccountManager: tsAccountManager,
)
private lazy var listMediaManager = BackupListMediaManagerImpl(
accountKeyStore: accountKeyStore,
attachmentStore: attachmentStore,
attachmentUploadStore: attachmentUploadStore,
backupAttachmentDownloadProgress: backupAttachmentDownloadProgress,
backupAttachmentDownloadStore: backupAttachmentDownloadStore,
backupAttachmentUploadProgress: backupAttachmentUploadProgress,
backupAttachmentUploadScheduler: backupAttachmentUploadScheduler,
backupAttachmentUploadStore: backupAttachmentUploadStore,
backupAttachmentUploadEraStore: backupAttachmentUploadEraStore,
backupListMediaStore: backupListMediaStore,
backupMediaErrorNotificationPresenter: backupMediaErrorNotificationPresenter,
backupRequestManager: backupRequestManager,
backupSettingsStore: backupSettingsStore,
dateProvider: dateProvider,
db: db,
notificationPresenter: notificationPresenter,
orphanedBackupAttachmentStore: orphanedBackupAttachmentStore,
remoteConfigManager: remoteConfigManager,
tsAccountManager: tsAccountManager,
)
@Test
func testListMedia() async throws {
let localUploadEra = "1"
let remoteConfigCdnNumber: UInt32 = 100
remoteConfigManager._currentConfig = RemoteConfig(
clockSkew: 0,
valueFlags: ["global.backups.mediaTierFallbackCdnNumber": "\(remoteConfigCdnNumber)"],
)
let mediaRootBackupKey = MediaRootBackupKey(backupKey: .generateRandom())
await db.awaitableWrite { tx in
accountKeyStore.setMediaRootBackupKey(mediaRootBackupKey, tx: tx)
backupSettingsStore.setBackupPlan(.paid(optimizeLocalStorage: false), tx: tx)
}
// Make a few attachments so we hit a couple pages of results
// from the request and from local db reads.
let numAttachmentsPerCase = 50
// There are N cases we care about:
// Case 1: Attachment exists locally but not on CDN
let localOnlyIds = await db.awaitableWrite { tx in
return (0..<numAttachmentsPerCase).map { _ in
return insertAttachment(
mediaName: UUID().uuidString,
mediaTierInfo: .init(
cdnNumber: 1,
unencryptedByteCount: 100,
sha256ContentHash: UUID().data,
incrementalMacInfo: nil,
uploadEra: localUploadEra,
lastDownloadAttemptTimestamp: nil,
),
tx: tx,
)
}
}
// Case 2: Attachment exists on cdn but not locally
let orphanCdnNumber: UInt32 = 99
let remoteOnlyCdnNumberMedia = (0..<numAttachmentsPerCase).map { _ in
return BackupArchive.Response.StoredMedia(
cdn: orphanCdnNumber,
mediaId: UUID().uuidString,
objectLength: 100,
)
}
// For other cases, we'll add duplicate entries on cdn at different
// cdn numbers that should be orphaned like case 2
var orphanCdnNumberMedia = [BackupArchive.Response.StoredMedia]()
// Case 3: exists on both, local didn't have cdn info
let discoveredCdnNumber: UInt32 = 5
var discoveredCdnNumberMedia = [BackupArchive.Response.StoredMedia]()
let discoveredCdnNumberIds = try await db.awaitableWrite { tx in
return try (0..<numAttachmentsPerCase).map { _ in
let mediaName = UUID().uuidString
let mediaId = try mediaRootBackupKey.deriveMediaId(mediaName)
discoveredCdnNumberMedia.append(.init(
cdn: discoveredCdnNumber,
mediaId: mediaId.asBase64Url,
objectLength: 100,
))
return insertAttachment(
mediaName: mediaName,
mediaTierInfo: nil,
tx: tx,
)
}
}
// Case 4: exists on both, cdn number matches.
let matchingCdnNumber: UInt32 = 3
var matchingCdnNumberMedia = [BackupArchive.Response.StoredMedia]()
let matchingCdnNumberIds = try await db.awaitableWrite { tx in
return try (0..<numAttachmentsPerCase).map { _ in
let mediaName = UUID().uuidString
let mediaId = try mediaRootBackupKey.deriveMediaId(mediaName)
matchingCdnNumberMedia.append(.init(
cdn: matchingCdnNumber,
mediaId: mediaId.asBase64Url,
objectLength: 100,
))
orphanCdnNumberMedia.append(.init(
cdn: orphanCdnNumber,
mediaId: mediaId.asBase64Url,
objectLength: 100,
))
return insertAttachment(
mediaName: mediaName,
mediaTierInfo: .init(
cdnNumber: matchingCdnNumber,
unencryptedByteCount: 100,
sha256ContentHash: UUID().data,
incrementalMacInfo: nil,
uploadEra: localUploadEra,
lastDownloadAttemptTimestamp: nil,
),
tx: tx,
)
}
}
// Case 5: exists on both, cdn number doesn't match.
var nonMatchingCdnNumberMedia = [BackupArchive.Response.StoredMedia]()
let nonMatchingCdnNumberIds = try await db.awaitableWrite { tx in
return try (0..<numAttachmentsPerCase).map { _ in
let mediaName = UUID().uuidString
let mediaId = try mediaRootBackupKey.deriveMediaId(mediaName)
nonMatchingCdnNumberMedia.append(.init(
// Prefer a cdn number matching remote config,
// instead of the other orphaned one below
cdn: remoteConfigCdnNumber,
mediaId: mediaId.asBase64Url,
objectLength: 100,
))
orphanCdnNumberMedia.append(.init(
cdn: remoteConfigCdnNumber,
mediaId: mediaId.asBase64Url,
objectLength: 100,
))
return insertAttachment(
mediaName: mediaName,
mediaTierInfo: .init(
cdnNumber: matchingCdnNumber,
unencryptedByteCount: 100,
sha256ContentHash: UUID().data,
incrementalMacInfo: nil,
uploadEra: localUploadEra,
lastDownloadAttemptTimestamp: nil,
),
tx: tx,
)
}
}
// Set up mock list response
backupRequestManager.listMediaResults = [
BackupArchive.Response.ListMediaResult(
storedMediaObjects: remoteOnlyCdnNumberMedia,
backupDir: "",
mediaDir: "",
cursor: "someCursor",
),
BackupArchive.Response.ListMediaResult(
storedMediaObjects: discoveredCdnNumberMedia,
backupDir: "",
mediaDir: "",
cursor: "someCursor",
),
BackupArchive.Response.ListMediaResult(
storedMediaObjects: matchingCdnNumberMedia,
backupDir: "",
mediaDir: "",
cursor: "someCursor",
),
BackupArchive.Response.ListMediaResult(
storedMediaObjects: nonMatchingCdnNumberMedia,
backupDir: "",
mediaDir: "",
cursor: "someCursor",
),
BackupArchive.Response.ListMediaResult(
storedMediaObjects: orphanCdnNumberMedia,
backupDir: "",
mediaDir: "",
cursor: nil,
),
]
try await listMediaManager.queryListMediaIfNeeded()
// Case 1 should've been marked as not uploaded to media tier,
// removed from download queue and added to upload queue.
db.read { tx in
let enqueuedAttachmentIds = backupAttachmentUploadStore.fetchNextUploads(
count: 60,
isFullsize: true,
tx: tx,
).map(\.attachmentRowId)
for attachmentId in localOnlyIds {
let attachment = attachmentStore.fetch(id: attachmentId, tx: tx)!
#expect(attachment.mediaTierInfo == nil)
#expect(backupAttachmentDownloadStore.getEnqueuedDownload(
attachmentRowId: attachmentId,
thumbnail: false,
tx: tx,
) == nil)
#expect(enqueuedAttachmentIds.contains(attachmentId))
}
}
// Case 2, and other duplicates, should've been marked for deletion.
db.read { tx in
for orphanMedia in remoteOnlyCdnNumberMedia + orphanCdnNumberMedia {
let mediaId = try! Data.data(fromBase64Url: orphanMedia.mediaId)
#expect(
try! OrphanedBackupAttachment
.filter(Column(OrphanedBackupAttachment.CodingKeys.mediaId) == mediaId)
.filter(Column(OrphanedBackupAttachment.CodingKeys.cdnNumber) == orphanMedia.cdn)
.fetchCount(tx.database)
== 1,
)
}
}
// Case 3 should be updated with cdn info, with uploads dequeued.
db.read { tx in
for attachmentId in discoveredCdnNumberIds {
let attachment = attachmentStore.fetch(id: attachmentId, tx: tx)!
#expect(attachment.mediaTierInfo?.cdnNumber == discoveredCdnNumber)
#expect(
try! QueuedBackupAttachmentUpload
.filter(Column(QueuedBackupAttachmentUpload.CodingKeys.attachmentRowId) == attachmentId)
.fetchCount(tx.database)
== 0,
)
}
}
// Case 4 should be untouched
db.read { tx in
for attachmentId in matchingCdnNumberIds {
let attachment = attachmentStore.fetch(id: attachmentId, tx: tx)!
#expect(attachment.mediaTierInfo?.cdnNumber == matchingCdnNumber)
}
}
// Case 5 should be updated with the remote cdn number
db.read { tx in
for attachmentId in nonMatchingCdnNumberIds {
let attachment = attachmentStore.fetch(id: attachmentId, tx: tx)!
#expect(attachment.mediaTierInfo?.cdnNumber == remoteConfigCdnNumber)
}
}
}
// MARK: - Helpers
typealias Attachment = SignalServiceKit.Attachment
private func insertAttachment(
mediaName: String,
mediaTierInfo: Attachment.MediaTierInfo?,
tx: DBWriteTransaction,
) -> Attachment.IDType {
let thread = TSThread(uniqueId: UUID().uuidString)
try! thread.insert(tx.database)
var attachmentRecord = Attachment.Record.mockStream(mediaName: mediaName)
try! attachmentRecord.insert(tx.database)
if let mediaTierInfo {
let attachment = try! Attachment(record: attachmentRecord)
attachment.mediaTierInfo = mediaTierInfo
attachmentRecord = Attachment.Record(attachment: attachment)
try! attachmentRecord.update(tx.database)
}
// We make all the attachments just thread wallpapers for ease of setup;
// for list media purposes it doesn't matter if its a message attachment
// or thread wallpaper attachment.
let referenceParams = AttachmentReference.ConstructionParams.mock(
owner: .thread(.threadWallpaperImage(.init(
threadRowId: thread.sqliteRowId!,
creationTimestamp: 0,
))),
)
_ = attachmentStore.addReference(
referenceParams,
attachmentRowId: attachmentRecord.sqliteId!,
tx: tx,
)
return attachmentRecord.sqliteId!
}
}