| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434 |
- // Copyright 2021 Google LLC
- //
- // Licensed under the Apache License, Version 2.0 (the "License");
- // you may not use this file except in compliance with the License.
- // You may obtain a copy of the License at
- //
- // http://www.apache.org/licenses/LICENSE-2.0
- //
- // Unless required by applicable law or agreed to in writing, software
- // distributed under the License is distributed on an "AS IS" BASIS,
- // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- // See the License for the specific language governing permissions and
- // limitations under the License.
- // Skip keychain tests on Catalyst and macOS. Tests are skipped because they
- // involve interactions with the keychain that require a provisioning profile.
- // This is because MLModelDownloader depends on Installations.
- // See go/firebase-macos-keychain-popups for more details.
- #if !targetEnvironment(macCatalyst) && !os(macOS)
- import XCTest
- @testable import FirebaseCore
- @testable import FirebaseInstallations
- @testable import FirebaseMLModelDownloader
- /// Mock options to configure default Firebase app.
- private enum MockOptions {
- static let appID = "1:123:ios:123abc"
- static let gcmSenderID = "mock-sender-id"
- static let projectID = "mock-project-id"
- static let apiKey = "ABcdEf-APIKeyWithValidFormat_0123456789"
- }
- // TODO: Create separate files for each class tested.
- final class ModelDownloaderUnitTests: XCTestCase {
- let fakeModelName = "fakeModelName"
- let fakemodelPath = "fakemodelPath"
- let fakeModelHash = "fakeModelHash"
- let fakeDownloadURL = URL(string: "www.fake-download-url.com")!
- let fakeFileURL = URL(string: "www.fake-model-file.com")!
- let fakeExpiryTime = "2021-01-20T04:20:10.220Z"
- let fakeModelSize = 20
- let fakeProjectID = "fakeProjectID"
- let fakeAPIKey = "fakeAPIKey"
- let fakeAppName = "fakeAppName"
- var fakeModelJSON: String {
- """
- {
- "downloadUri":"\(fakeDownloadURL)",
- "expireTime":"\(fakeExpiryTime)",
- "sizeBytes":"\(fakeModelSize)"
- }
- """
- }
- let successAuthTokenProvider =
- { (completion: @escaping (Result<String, DownloadError>) -> Void) in
- completion(.success("fakeFISToken"))
- }
- override class func setUp() {
- let options = FirebaseOptions(
- googleAppID: MockOptions.appID,
- gcmSenderID: MockOptions.gcmSenderID
- )
- options.apiKey = MockOptions.apiKey
- options.projectID = MockOptions.projectID
- FirebaseApp.configure(options: options)
- FirebaseConfiguration.shared.setLoggerLevel(.debug)
- }
- override func setUp() {
- do {
- try ModelFileManager.emptyModelsDirectory()
- } catch {
- XCTFail("Could not empty models directory.")
- }
- }
- override class func tearDown() {
- try? FileManager.default.removeItem(atPath: NSTemporaryDirectory())
- FirebaseApp.app()?.delete { _ in }
- }
- /// Test if the same downloader instance is returned for an app.
- func testModelDownloaderInit() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader1 = ModelDownloader.modelDownloader(app: testApp)
- let modelDownloader2 = ModelDownloader.modelDownloader()
- XCTAssert(modelDownloader1 === modelDownloader2)
- }
- /// Test if app deletion removes associated downloader instance.
- func testAppDeletion() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader1 = ModelDownloader.modelDownloader()
- testApp.delete { success in
- XCTAssertTrue(success)
- }
- ModelDownloaderUnitTests.setUp()
- let modelDownloader2 = ModelDownloader.modelDownloader()
- XCTAssert(modelDownloader1 !== modelDownloader2)
- }
- /// Test invalid model file deletion.
- func testDeleteInvalidModel() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- modelDownloader.deleteDownloadedModel(name: fakeModelName) { result in
- switch result {
- case .success: XCTFail("Unexpected successful model deletion.")
- case let .failure(error):
- XCTAssertNotNil(error)
- }
- }
- }
- func testGetModelNameFromFilePath() {
- var fakeURL = URL(fileURLWithPath: "modelDirectory/@@appName@@\(fakeModelName).tflite")
- XCTAssertEqual(ModelFileManager.getModelNameFromFilePath(fakeURL), fakeModelName)
- fakeURL = URL(fileURLWithPath: "modelDirectory/--appName--\(fakeModelName).tflite")
- XCTAssertNil(ModelFileManager.getModelNameFromFilePath(fakeURL))
- }
- /// Test listing models.
- func testListModels() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let testName = String(#function.dropLast(2))
- let modelDownloader = ModelDownloader.modelDownloaderWithDefaults(
- .createUnitTestInstance(testName: testName),
- app: testApp
- )
- let tempModelFileURL = tempModelFile()
- let completionExpectation = expectation(description: "Completion handler")
- let localModelInfo = fakeLocalModelInfo()
- localModelInfo.writeToDefaults(
- .getUnitTestInstance(testName: testName),
- appName: testApp.name
- )
- modelDownloader.listDownloadedModels { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(models):
- /// Expected success since model info was written to user defaults.
- XCTAssertEqual(models.first?.path, tempModelFileURL.path)
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempModelFileURL)
- }
- /// Test invalid path in model directory.
- func testListModelsInvalidPath() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let invalidTempModelFileURL = tempModelFile(invalid: true)
- let completionExpectation = expectation(description: "Completion handler")
- modelDownloader.listDownloadedModels { result in
- completionExpectation.fulfill()
- switch result {
- case .success:
- XCTFail("Unexpected success of list models.")
- case let .failure(error):
- switch error {
- case let .internalError(description):
- XCTAssertTrue(description
- .contains("List models failed due to unexpected model file name"))
- default: XCTFail("Expected failure due to bad model name.")
- }
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: invalidTempModelFileURL)
- }
- /// Test listing models without local info in user defaults.
- func testListModelsNoLocalInfo() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let testName = String(#function.dropLast(2))
- let tempModelFileURL = tempModelFile()
- let modelDownloader = ModelDownloader.modelDownloaderWithDefaults(
- .createUnitTestInstance(testName: testName),
- app: testApp
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelDownloader.listDownloadedModels { result in
- completionExpectation.fulfill()
- switch result {
- case .success:
- XCTFail("Unexpected success of list models.")
- case let .failure(error):
- switch error {
- /// Expected failure since no model info was written to user defaults.
- case let .internalError(description):
- XCTAssertTrue(description.contains("List models failed due to no local model info"))
- default: XCTFail("Expected failure due to bad model name.")
- }
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempModelFileURL)
- }
- /// Test how URL conversion behaves if there are spaces in the path.
- func testURLConversion() {
- /// Spaces in the string only convert to URL when using URL(fileURLWithPath: ).
- let fakeURLWithSpace = URL(string: "file:///fakeDir1/fake%20Dir2/fakeFile")!
- XCTAssertEqual(fakeURLWithSpace, URL(string: fakeURLWithSpace.absoluteString))
- XCTAssertEqual(fakeURLWithSpace, URL(fileURLWithPath: fakeURLWithSpace.path))
- XCTAssertNil(URL(string: fakeURLWithSpace.path))
- /// Strings without spaces should work fine either way.
- let fakeURLWithoutSpace = URL(string: "fakeDir1/fakeDir2/fakeFile")!
- XCTAssertEqual(fakeURLWithoutSpace, URL(string: fakeURLWithoutSpace.absoluteString)!)
- XCTAssertEqual(fakeURLWithoutSpace, URL(string: fakeURLWithoutSpace.path)!)
- }
- /// Test on-device logging.
- func testDeviceLogging() {
- DeviceLogger.logEvent(level: .error,
- message: "This is a test error logged to device.",
- messageCode: .testError)
- }
- /// Compare proto serialization methods.
- func testTelemetryEncoding() {
- let fakeModel = CustomModel(
- name: fakeModelName,
- size: 10,
- path: fakemodelPath,
- hash: fakeModelHash
- )
- var modelOptions = ModelOptions()
- modelOptions.setModelOptions(modelName: fakeModel.name, modelHash: fakeModel.hash)
- guard let binaryData = try? modelOptions.serializedData(),
- let jsonData = try? modelOptions.jsonUTF8Data(),
- let binaryEvent = try? ModelOptions(serializedData: binaryData),
- let jsonEvent = try? ModelOptions(jsonUTF8Data: jsonData) else {
- XCTFail("Encoding error.")
- return
- }
- XCTAssertNotNil(binaryData)
- XCTAssertNotNil(jsonData)
- XCTAssertLessThan(binaryData.count, jsonData.count)
- XCTAssertEqual(binaryEvent, jsonEvent)
- }
- /// Get model info if server returns a new model info.
- func testGetModelInfoWith200() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case let .modelInfo(remoteModelInfo):
- XCTAssertEqual(remoteModelInfo.name, self.fakeModelName)
- XCTAssertEqual(remoteModelInfo.downloadURL, self.fakeDownloadURL)
- XCTAssertEqual(remoteModelInfo.size, self.fakeModelSize)
- XCTAssertEqual(remoteModelInfo.modelHash, self.fakeModelHash)
- case .notModified: XCTFail("Expected new model info from server.")
- }
- case let .failure(error): XCTFail("Failed with error - \(error)")
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model info if model info is not modified.
- func testGetModelInfoWith304() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 304)
- let localModelInfo = fakeLocalModelInfo(mismatch: false)
- let modelInfoRetriever = fakeModelRetriever(
- fakeSession: session,
- fakeLocalModelInfo: localModelInfo
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case .modelInfo: XCTFail("Unexpected new model info from server.")
- case .notModified: break
- }
- case let .failure(error): XCTFail("Failed with error - \(error)")
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model info if model info is not modified but local model info is not set.
- func testGetModelInfoWith304MissingLocalInfo() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 304)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case .modelInfo: XCTFail("Unexpected new model info from server.")
- case .notModified: XCTFail("Expected failure since local model info was not set.")
- }
- case let .failure(error):
- switch error {
- case let .internalError(description: errorMessage): XCTAssertEqual(errorMessage,
- "Model info was deleted unexpectedly.")
- default: XCTFail("Expected failure since local model info was not set.")
- }
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model info if model info is not modified but local model info is not set.
- func testGetModelInfoWith304HashMismatch() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 304)
- let localModelInfo = fakeLocalModelInfo(mismatch: true)
- let modelInfoRetriever = fakeModelRetriever(
- fakeSession: session,
- fakeLocalModelInfo: localModelInfo
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case .modelInfo: XCTFail("Unexpected new model info from server.")
- case .notModified: XCTFail("Expected failure since model hash does not match.")
- }
- case let .failure(error):
- switch error {
- case let .internalError(description: errorMessage): XCTAssertEqual(errorMessage,
- "Unexpected model hash value.")
- default: XCTFail("Expected failure since local model info was not set.")
- }
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model info if model name is invalid.
- func testGetModelInfoWith400() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 400)
- let localModelInfo = fakeLocalModelInfo(mismatch: false)
- let modelInfoRetriever = fakeModelRetriever(
- fakeSession: session,
- fakeLocalModelInfo: localModelInfo
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case .modelInfo: XCTFail("Unexpected new model info from server.")
- case .notModified: XCTFail("Expected failure since model name is invalid.")
- }
- case let .failure(error):
- XCTAssertEqual(error, .invalidArgument)
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model info if model name does not exist.
- func testGetModelInfoWith404() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 404)
- let localModelInfo = fakeLocalModelInfo(mismatch: false)
- let modelInfoRetriever = fakeModelRetriever(
- fakeSession: session,
- fakeLocalModelInfo: localModelInfo
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case .modelInfo: XCTFail("Unexpected new model info from server.")
- case .notModified: XCTFail("Expected failure since model name is not availabl.")
- }
- case let .failure(error):
- XCTAssertEqual(error, .notFound)
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model info if resource exhausted due to too many requests.
- func testGetModelInfoWith429() {
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 429)
- let localModelInfo = fakeLocalModelInfo(mismatch: false)
- let modelInfoRetriever = fakeModelRetriever(
- fakeSession: session,
- fakeLocalModelInfo: localModelInfo
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(fakemodelInfoResult):
- switch fakemodelInfoResult {
- case .modelInfo: XCTFail("Unexpected new model info from server.")
- case .notModified: XCTFail("Expected failure since resource exhausted.")
- }
- case let .failure(error):
- XCTAssertEqual(error, .resourceExhausted)
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model if server returns a model file.
- func testModelDownloadWith200() {
- let tempFileURL = tempFile()
- let remoteModelInfo = fakeModelInfo(expired: false, oversized: false)
- let fakeResponse = HTTPURLResponse(url: fakeFileURL,
- statusCode: 200,
- httpVersion: nil,
- headerFields: nil)!
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let progressExpectation = expectation(description: "Progress handler")
- progressExpectation.expectedFulfillmentCount = 2
- let completionExpectation = expectation(description: "Completion handler")
- let taskProgressHandler: ModelDownloadTask.ProgressHandler = {
- progress in
- progressExpectation.fulfill()
- XCTAssertEqual(progress, 0.4)
- }
- let taskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(model):
- self.checkModel(model: model)
- self.deleteFile(at: model.path)
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- let testName = String(#function.dropLast(2))
- let modelDownloadTask = ModelDownloadTask(remoteModelInfo: remoteModelInfo,
- appName: "fakeAppName",
- defaults: .createUnitTestInstance(
- testName: testName
- ),
- downloader: downloader,
- progressHandler: taskProgressHandler,
- completion: taskCompletion)
- modelDownloadTask.resume()
- wait(for: [fileDownloaderExpectation], timeout: 0.1)
- downloader.progressHandler?(100, 250)
- downloader.progressHandler?(100, 250)
- wait(for: [progressExpectation], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponse,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if server returns an error due to expired url.
- func testModelDownloadWith400Expired() {
- let tempFileURL = tempFile()
- let remoteModelInfo = fakeModelInfo(expired: true, oversized: false)
- let fakeResponse = HTTPURLResponse(url: fakeFileURL,
- statusCode: 400,
- httpVersion: nil,
- headerFields: nil)!
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let completionExpectation = expectation(description: "Completion handler")
- let taskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- XCTAssertEqual(error, .expiredDownloadURL)
- }
- }
- let testName = String(#function.dropLast(2))
- let modelDownloadTask = ModelDownloadTask(remoteModelInfo: remoteModelInfo,
- appName: "fakeAppName",
- defaults: .createUnitTestInstance(
- testName: testName
- ),
- downloader: downloader,
- completion: taskCompletion)
- modelDownloadTask.resume()
- wait(for: [fileDownloaderExpectation], timeout: 0.1)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponse,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if server returns an error due to invalid model name.
- func testModelDownloadWith400Invalid() {
- let tempFileURL = tempFile()
- let remoteModelInfo = fakeModelInfo(expired: false, oversized: false)
- let fakeResponse = HTTPURLResponse(url: fakeFileURL,
- statusCode: 400,
- httpVersion: nil,
- headerFields: nil)!
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let completionExpectation = expectation(description: "Completion handler")
- let taskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- XCTAssertEqual(error, .invalidArgument)
- }
- }
- let testName = String(#function.dropLast(2))
- let modelDownloadTask = ModelDownloadTask(remoteModelInfo: remoteModelInfo,
- appName: "fakeAppName",
- defaults: .createUnitTestInstance(
- testName: testName
- ),
- downloader: downloader,
- completion: taskCompletion)
- modelDownloadTask.resume()
- wait(for: [fileDownloaderExpectation], timeout: 0.1)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponse,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model info if Firebase ML API is not enabled or permission denied.
- func testGetModelInfoWith403() {
- let errorMessage =
- """
- {
- "error":
- {
- "message":"API Disabled."
- }
- }
- """
- let session = fakeModelInfoSessionWithURL(
- fakeDownloadURL,
- statusCode: 403,
- customJSON: errorMessage
- )
- let localModelInfo = fakeLocalModelInfo(mismatch: false)
- let modelInfoRetriever = fakeModelRetriever(
- fakeSession: session,
- fakeLocalModelInfo: localModelInfo
- )
- let completionExpectation = expectation(description: "Completion handler")
- modelInfoRetriever.downloadModelInfo { result in
- completionExpectation.fulfill()
- switch result {
- case .success:
- XCTFail("Unexpected new model info from server.")
- case let .failure(error):
- XCTAssertEqual(error, .permissionDenied)
- }
- }
- wait(for: [completionExpectation], timeout: 0.5)
- }
- /// Get model if multiple duplicate requests are made.
- func testModelDownloadWithMergeRequests() {
- let tempFileURL = tempFile()
- let numberOfRequests = 5
- let remoteModelInfo = fakeModelInfo(expired: false, oversized: false)
- let fakeResponse = HTTPURLResponse(url: fakeFileURL,
- statusCode: 200,
- httpVersion: nil,
- headerFields: nil)!
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let completionExpectation = expectation(description: "Completion handler")
- completionExpectation.expectedFulfillmentCount = numberOfRequests + 1
- let taskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(model):
- self.checkModel(model: model)
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- let testName = String(#function.dropLast(2))
- let modelDownloadTask = ModelDownloadTask(remoteModelInfo: remoteModelInfo,
- appName: fakeAppName,
- defaults: .createUnitTestInstance(
- testName: testName
- ),
- downloader: downloader,
- completion: taskCompletion)
- for i in 1 ... numberOfRequests {
- let newTaskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(model):
- self.checkModel(model: model)
- if i == numberOfRequests {
- self.deleteFile(at: model.path)
- }
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- modelDownloadTask.resume()
- modelDownloadTask.merge(newCompletion: newTaskCompletion)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponse,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation], timeout: 1)
- wait(for: [completionExpectation], timeout: 2)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if multiple duplicate requests are made (test at the API surface).
- func testGetModelWithMergeRequests() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let tempFileURL = tempFile()
- let numberOfRequests = 6
- let fakeResponse = HTTPURLResponse(url: fakeFileURL,
- statusCode: 200,
- httpVersion: nil,
- headerFields: nil)!
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let completionExpectation = expectation(description: "Completion handler")
- completionExpectation.expectedFulfillmentCount = numberOfRequests
- for i in 1 ... numberOfRequests {
- modelDownloader.downloadInfoAndModel(modelName: "fakeModelName",
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(model):
- self.checkModel(model: model)
- if i == numberOfRequests {
- self.deleteFile(at: model.path)
- }
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- }
- wait(for: [fileDownloaderExpectation], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponse,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 1)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if server returns an error due to model not found.
- func testModelDownloadWith404() {
- let tempFileURL = tempFile()
- let remoteModelInfo = fakeModelInfo(expired: false, oversized: false)
- let fakeResponse = HTTPURLResponse(url: fakeFileURL,
- statusCode: 404,
- httpVersion: nil,
- headerFields: nil)!
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let completionExpectation = expectation(description: "Completion handler")
- let taskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- XCTAssertEqual(error, .notFound)
- }
- }
- let testName = String(#function.dropLast(2))
- let modelDownloadTask = ModelDownloadTask(remoteModelInfo: remoteModelInfo,
- appName: "fakeAppName",
- defaults: .createUnitTestInstance(
- testName: testName
- ),
- downloader: downloader,
- completion: taskCompletion)
- modelDownloadTask.resume()
- wait(for: [fileDownloaderExpectation], timeout: 0.1)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponse,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if server returns an error due to model not found.
- func testModelDownloadNetworkError() {
- let tempFileURL = tempFile()
- let remoteModelInfo = fakeModelInfo(expired: false, oversized: false)
- let downloader = MockModelFileDownloader()
- let fileDownloaderExpectation = expectation(description: "File downloader")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- let completionExpectation = expectation(description: "Completion handler")
- let taskCompletion: ModelDownloadTask.Completion = { result in
- completionExpectation.fulfill()
- switch result {
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- XCTAssertEqual(error, .failedPrecondition)
- }
- }
- let testName = String(#function.dropLast(2))
- let modelDownloadTask = ModelDownloadTask(remoteModelInfo: remoteModelInfo,
- appName: "fakeAppName",
- defaults: .createUnitTestInstance(
- testName: testName
- ),
- downloader: downloader,
- completion: taskCompletion)
- modelDownloadTask.resume()
- wait(for: [fileDownloaderExpectation], timeout: 0.1)
- let offlineError = DownloadError.internalError(description: "Network appears to be offline.")
- downloader
- .completion?(.failure(FileDownloaderError.networkError(offlineError)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model with a valid server response.
- func testGetModelSuccessful() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseNotFound = HTTPURLResponse(url: fakeFileURL,
- statusCode: 200,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation = expectation(description: "File downloader")
- let completionExpectation = expectation(description: "Completion handler")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- case let .success(model):
- self.checkModel(model: model)
- self.deleteFile(at: model.path)
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- wait(for: [fileDownloaderExpectation], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseNotFound,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if download url expired but retry was successful.
- func testGetModelSuccessfulRetry() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseExpired = HTTPURLResponse(url: fakeFileURL,
- statusCode: 400,
- httpVersion: nil,
- headerFields: nil)!
- let fakeResponseSucceeded = HTTPURLResponse(url: fakeFileURL,
- statusCode: 200,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation1 = expectation(description: "File downloader before retry")
- let completionExpectation = expectation(description: "Completion handler.")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation1.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- // Retry successful - model returned.
- case let .success(model):
- self.checkModel(model: model)
- self.deleteFile(at: model.path)
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- wait(for: [fileDownloaderExpectation1], timeout: 0.5)
- // Handles response after one retry.
- let fileDownloaderExpectation2 = expectation(description: "File downloader after retry")
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation2.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation2], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseSucceeded,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if model doesn't exist.
- func testGetModelNotFound() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseNotFound = HTTPURLResponse(url: fakeFileURL,
- statusCode: 404,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation = expectation(description: "File downloader before retry")
- let completionExpectation = expectation(description: "Completion handler")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- // Retry failed due to another expired url - error returned.
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- switch error {
- case .notFound: break
- default: XCTFail("Expected failure due to model not found.")
- }
- }
- }
- wait(for: [fileDownloaderExpectation], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseNotFound,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if invalid or insufficient permissions.
- func testGetModelUnauthorized() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseNotFound = HTTPURLResponse(url: fakeFileURL,
- statusCode: 403,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation = expectation(description: "File downloader")
- let completionExpectation = expectation(description: "Completion handler")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- switch error {
- case .permissionDenied: break
- default: XCTFail("Expected failure due to bad permissions.")
- }
- }
- }
- wait(for: [fileDownloaderExpectation], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseNotFound,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if download url expired and retry url also expired.
- func testGetModelFailedRetry() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseExpired = HTTPURLResponse(url: fakeFileURL,
- statusCode: 400,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation1 = expectation(description: "File downloader before retry")
- let completionExpectation = expectation(description: "Completion handler")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation1.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- // Retry failed due to another expired url - error returned.
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- switch error {
- case let .internalError(description: errorMessage): XCTAssertEqual(errorMessage,
- "Unable to update expired model info.")
- default: XCTFail("Expected failure due to expired model info.")
- }
- }
- }
- wait(for: [fileDownloaderExpectation1], timeout: 0.5)
- let fileDownloaderExpectation2 = expectation(description: "File downloader after retry")
- // Handles response after one retry.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation2.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation2], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if multiple retries all returned expired urls.
- func testGetModelMultipleFailedRetries() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- modelDownloader.numberOfRetries = 2
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseExpired = HTTPURLResponse(url: fakeFileURL,
- statusCode: 400,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation1 = expectation(description: "File downloader before retry")
- let completionExpectation = expectation(description: "Completion handler")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation1.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- // Two retries failed due to expired urls - error returned.
- case .success: XCTFail("Unexpected successful model download.")
- case let .failure(error):
- switch error {
- case let .internalError(description: errorMessage): XCTAssertEqual(errorMessage,
- "Unable to update expired model info.")
- default: XCTFail("Expected failure due to expired model info.")
- }
- }
- }
- wait(for: [fileDownloaderExpectation1], timeout: 0.5)
- let fileDownloaderExpectation2 = expectation(description: "file downloader after first retry")
- // Handles response after first retry.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation2.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation2], timeout: 0.5)
- let fileDownloaderExpectation3 =
- expectation(description: "File downloader after second retry")
- // Handles response after second retry.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation3.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation3], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- /// Get model if a retry finally returns an unexpired download url.
- func testGetModelMultipleRetriesWithSuccess() {
- guard let testApp = FirebaseApp.app() else {
- XCTFail("Default app was not configured.")
- return
- }
- testApp.isDataCollectionDefaultEnabled = false
- let modelDownloader = ModelDownloader.modelDownloader()
- modelDownloader.numberOfRetries = 4
- let conditions = ModelDownloadConditions()
- let session = fakeModelInfoSessionWithURL(fakeDownloadURL, statusCode: 200)
- let modelInfoRetriever = fakeModelRetriever(fakeSession: session)
- let downloader = MockModelFileDownloader()
- let tempFileURL = tempFile()
- let fakeResponseExpired = HTTPURLResponse(url: fakeFileURL,
- statusCode: 400,
- httpVersion: nil,
- headerFields: nil)!
- let fakeResponseSucceeded = HTTPURLResponse(url: fakeFileURL,
- statusCode: 200,
- httpVersion: nil,
- headerFields: nil)!
- let fileDownloaderExpectation1 = expectation(description: "File downloader before retry")
- let completionExpectation = expectation(description: "Completion handler")
- // Handles initial response.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation1.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- modelDownloader.downloadInfoAndModel(modelName: fakeModelName,
- modelInfoRetriever: modelInfoRetriever,
- downloader: downloader,
- conditions: conditions) { result in
- completionExpectation.fulfill()
- switch result {
- // One retry failed but the next retry succeeded - model returned.
- case let .success(model):
- self.checkModel(model: model)
- self.deleteFile(at: model.path)
- case let .failure(error):
- XCTFail("Error - \(error)")
- }
- }
- wait(for: [fileDownloaderExpectation1], timeout: 0.5)
- let fileDownloaderExpectation2 = expectation(description: "File downloader after first retry")
- // Handles response after first retry.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation2.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation2], timeout: 0.5)
- let fileDownloaderExpectation3 =
- expectation(description: "File downloader after second retry")
- // Handles response after second retry.
- downloader.downloadFileHandler = { url in
- fileDownloaderExpectation3.fulfill()
- XCTAssertEqual(url, self.fakeDownloadURL)
- }
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseExpired,
- fileURL: tempFileURL)))
- wait(for: [fileDownloaderExpectation3], timeout: 0.5)
- downloader
- .completion?(.success(FileDownloaderResponse(urlResponse: fakeResponseSucceeded,
- fileURL: tempFileURL)))
- wait(for: [completionExpectation], timeout: 0.5)
- try? ModelFileManager.removeFile(at: tempFileURL)
- }
- }
- /// Mock URL session for testing.
- class MockModelInfoRetrieverSession: ModelInfoRetrieverSession {
- var data: Data?
- var response: URLResponse?
- var error: Error?
- func getModelInfo(with request: URLRequest,
- completion: @escaping (Data?, URLResponse?, Error?) -> Void) {
- completion(data, response, error)
- }
- }
- class MockModelFileDownloader: FileDownloader {
- var progressHandler: ProgressHandler?
- var completion: CompletionHandler?
- var downloadFileHandler: ((_ url: URL) -> Void)?
- func downloadFile(with url: URL, progressHandler: @escaping ProgressHandler,
- completion: @escaping CompletionHandler) {
- self.progressHandler = progressHandler
- self.completion = completion
- downloadFileHandler?(url)
- }
- }
- extension UserDefaults {
- /// Returns a new cleared instance of user defaults.
- static func createUnitTestInstance(testName: String) -> UserDefaults {
- let suiteName = "com.google.firebase.ml.test.\(testName)"
- let defaults = UserDefaults(suiteName: suiteName)!
- defaults.removePersistentDomain(forName: suiteName)
- return defaults
- }
- /// Returns the existing user defaults instance.
- static func getUnitTestInstance(testName: String) -> UserDefaults {
- let suiteName = "com.google.firebase.ml.test.\(testName)"
- return UserDefaults(suiteName: suiteName)!
- }
- }
- // MARK: - Helpers
- extension ModelDownloaderUnitTests {
- func fakeModelInfoSessionWithURL(_ url: URL, statusCode: Int,
- customJSON: String? = nil) -> MockModelInfoRetrieverSession {
- let fakeSession = MockModelInfoRetrieverSession()
- let fakeJSON = customJSON ?? fakeModelJSON
- fakeSession.data = fakeJSON.data(using: .utf8)
- fakeSession.response = HTTPURLResponse(
- url: url,
- statusCode: statusCode,
- httpVersion: nil,
- headerFields: ["Etag": "fakeModelHash"]
- )
- fakeSession.error = nil
- return fakeSession
- }
- func fakeModelRetriever(fakeSession: MockModelInfoRetrieverSession,
- fakeLocalModelInfo: LocalModelInfo? = nil) -> ModelInfoRetriever {
- // TODO: Replace with fake to check if it was used correctly by the download task.
- let modelInfoRetriever = ModelInfoRetriever(
- modelName: fakeModelName,
- projectID: fakeProjectID,
- apiKey: fakeAPIKey,
- appName: fakeAppName, authTokenProvider: successAuthTokenProvider,
- session: fakeSession, localModelInfo: fakeLocalModelInfo
- )
- return modelInfoRetriever
- }
- func checkModel(model: CustomModel) {
- XCTAssertEqual(model.name, fakeModelName)
- XCTAssertEqual(model.size, fakeModelSize)
- XCTAssertEqual(model.hash, fakeModelHash)
- let modelURL = URL(fileURLWithPath: model.path)
- XCTAssertTrue(ModelFileManager.isFileReachable(at: modelURL))
- }
- func tempFile(name: String = "fake-model-file.tmp") -> URL {
- let tempDirectoryURL = URL(fileURLWithPath: NSTemporaryDirectory(),
- isDirectory: true)
- let tempFileURL = tempDirectoryURL.appendingPathComponent(name)
- let tempData: Data = "fakeModelData".data(using: .utf8)!
- try? tempData.write(to: tempFileURL)
- return tempFileURL
- }
- func tempModelFile(invalid: Bool = false) -> URL {
- let modelDirectoryURL = ModelFileManager.modelsDirectory!
- let modelFileName = invalid ? "fbml_model_fake-model-file.tmp" :
- "fbml_model@@__FIRAPP_DEFAULT@@\(fakeModelName)"
- let tempFileURL = modelDirectoryURL.appendingPathComponent(modelFileName)
- .appendingPathExtension("tflite")
- let tempData: Data = "fakeModelData".data(using: .utf8)!
- try? tempData.write(to: tempFileURL)
- return tempFileURL
- }
- func deleteFile(at path: String) {
- let url = URL(fileURLWithPath: path)
- do {
- try ModelFileManager.removeFile(at: url)
- } catch {
- XCTFail("Failed to delete file at \(path).")
- }
- }
- func fakeLocalModelInfo(mismatch: Bool = false) -> LocalModelInfo {
- let hash = mismatch ? "wrongModelHash" : fakeModelHash
- return LocalModelInfo(
- name: fakeModelName,
- modelHash: hash,
- size: 20
- )
- }
- func fakeModelInfo(expired: Bool, oversized: Bool) -> RemoteModelInfo {
- let expiryTime = expired ? Date() : Date().addingTimeInterval(1000)
- let size = oversized ? Int(Int64.max) / 4 : fakeModelSize
- return RemoteModelInfo(name: fakeModelName,
- downloadURL: fakeDownloadURL,
- modelHash: fakeModelHash,
- size: size,
- urlExpiryTime: expiryTime)
- }
- }
- #endif // !targetEnvironment(macCatalyst) && !os(macOS)
|