existing mock files for template unit tests

Author: paulb777

FirebaseAI/Tests/Unit/TemplateChatSessionTests.swift

@@ -37,9 +37,9 @@ final class TemplateChatSessionTests: XCTestCase {
 
   func testSendMessage() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
-      forResource: "unary-success-response",
+      forResource: "unary-success-basic-reply-short",
       withExtension: "json",
-      subdirectory: "mock-responses",
+      subdirectory: "mock-responses/googleai",
       isTemplateRequest: true
     )
     let chat = model.startChat(template: "test-template")
@@ -48,6 +48,10 @@ final class TemplateChatSessionTests: XCTestCase {
     XCTAssertEqual(chat.history[0].role, "user")
     XCTAssertEqual((chat.history[0].parts.first as? TextPart)?.text, "Hello")
     XCTAssertEqual(chat.history[1].role, "model")
+    XCTAssertEqual(
+      (chat.history[1].parts.first as? TextPart)?.text,
+      "Google's headquarters, also known as the Googleplex, is located in **Mountain View, California**.\n"
+    )
     XCTAssertEqual(response.candidates.count, 1)
   }
 

FirebaseAI/Tests/Unit/TemplateGenerativeModelTests.swift

@@ -37,17 +37,20 @@ final class TemplateGenerativeModelTests: XCTestCase {
 
   func testGenerateContent() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
-      forResource: "unary-success-response",
+      forResource: "unary-success-basic-reply-short",
       withExtension: "json",
-      subdirectory: "mock-responses",
+      subdirectory: "mock-responses/googleai",
       isTemplateRequest: true
     )
 
     let response = try await model.generateContent(
       template: "test-template",
       variables: ["name": "test"]
     )
-    XCTAssertEqual(response.text, "Hello, world!")
+    XCTAssertEqual(
+      response.text,
+      "Google's headquarters, also known as the Googleplex, is located in **Mountain View, California**.\n"
+    )
   }
 
   func testGenerateContentStream() async throws {

FirebaseAI/Tests/Unit/TemplateImagenModelTests.swift

@@ -36,17 +36,17 @@ final class TemplateImagenModelTests: XCTestCase {
 
   func testGenerateImages() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
-      forResource: "unary-success-image-response",
+      forResource: "unary-success-generate-images-base64",
       withExtension: "json",
-      subdirectory: "mock-responses",
+      subdirectory: "mock-responses/vertexai",
       isTemplateRequest: true
     )
 
     let response = try await model.generateImages(
       template: "test-template",
       variables: ["prompt": "a cat picture"]
     )
-    XCTAssertEqual(response.images.count, 1)
-    XCTAssertEqual(response.images.first?.data, Data(base64Encoded: "aW1hZ2UgZGF0YQ=="))
+    XCTAssertEqual(response.images.count, 4)
+    XCTAssertNotNil(response.images.first?.data)
   }
 }