|  | 
|  | 1 | +// Copyright 2025 Google LLC | 
|  | 2 | +// | 
|  | 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); | 
|  | 4 | +// you may not use this file except in compliance with the License. | 
|  | 5 | +// You may obtain a copy of the License at | 
|  | 6 | +// | 
|  | 7 | +//      http://www.apache.org/licenses/LICENSE-2.0 | 
|  | 8 | +// | 
|  | 9 | +// Unless required by applicable law or agreed to in writing, software | 
|  | 10 | +// distributed under the License is distributed on an "AS IS" BASIS, | 
|  | 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | 
|  | 12 | +// See the License for the specific language governing permissions and | 
|  | 13 | +// limitations under the License. | 
|  | 14 | + | 
|  | 15 | +import XCTest | 
|  | 16 | + | 
|  | 17 | +@testable import FirebaseAILogic | 
|  | 18 | + | 
|  | 19 | +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) | 
|  | 20 | +@available(watchOS, unavailable) | 
|  | 21 | +final class BidiGenerateContentServerMessageTests: XCTestCase { | 
|  | 22 | +  let decoder = JSONDecoder() | 
|  | 23 | + | 
|  | 24 | +  func testDecodeBidiGenerateContentServerMessage_setupComplete() throws { | 
|  | 25 | +    let json = """ | 
|  | 26 | +    { | 
|  | 27 | +      "setupComplete" : {} | 
|  | 28 | +    } | 
|  | 29 | +    """ | 
|  | 30 | +    let jsonData = try XCTUnwrap(json.data(using: .utf8)) | 
|  | 31 | + | 
|  | 32 | +    let serverMessage = try decoder.decode(BidiGenerateContentServerMessage.self, from: jsonData) | 
|  | 33 | +    guard case .setupComplete = serverMessage.messageType else { | 
|  | 34 | +      XCTFail("Decoded message is not a setupComplete message.") | 
|  | 35 | +      return | 
|  | 36 | +    } | 
|  | 37 | +  } | 
|  | 38 | + | 
|  | 39 | +  func testDecodeBidiGenerateContentServerMessage_serverContent() throws { | 
|  | 40 | +    let json = """ | 
|  | 41 | +    { | 
|  | 42 | +      "serverContent" : { | 
|  | 43 | +        "modelTurn" : { | 
|  | 44 | +          "parts" : [ | 
|  | 45 | +            { | 
|  | 46 | +              "inlineData" : { | 
|  | 47 | +                "data" : "BQUFBQU=", | 
|  | 48 | +                "mimeType" : "audio/pcm" | 
|  | 49 | +              } | 
|  | 50 | +            } | 
|  | 51 | +          ], | 
|  | 52 | +          "role" : "model" | 
|  | 53 | +        }, | 
|  | 54 | +        "turnComplete": true, | 
|  | 55 | +        "groundingMetadata": { | 
|  | 56 | +          "webSearchQueries": ["query1", "query2"], | 
|  | 57 | +          "groundingChunks": [ | 
|  | 58 | +            { "web": { "uri": "uri1", "title": "title1" } } | 
|  | 59 | +          ], | 
|  | 60 | +          "groundingSupports": [ | 
|  | 61 | +            { "segment": { "endIndex": 10, "text": "text" }, "groundingChunkIndices": [0] } | 
|  | 62 | +          ], | 
|  | 63 | +          "searchEntryPoint": { "renderedContent": "html" } | 
|  | 64 | +        }, | 
|  | 65 | +        "inputTranscription": { | 
|  | 66 | +          "text": "What day of the week is it?" | 
|  | 67 | +        }, | 
|  | 68 | +        "outputTranscription": { | 
|  | 69 | +          "text": "Today is friday" | 
|  | 70 | +        } | 
|  | 71 | +      } | 
|  | 72 | +    } | 
|  | 73 | +    """ | 
|  | 74 | +    let jsonData = try XCTUnwrap(json.data(using: .utf8)) | 
|  | 75 | + | 
|  | 76 | +    let serverMessage = try decoder.decode(BidiGenerateContentServerMessage.self, from: jsonData) | 
|  | 77 | +    guard case let .serverContent(serverContent) = serverMessage.messageType else { | 
|  | 78 | +      XCTFail("Decoded message is not a serverContent message.") | 
|  | 79 | +      return | 
|  | 80 | +    } | 
|  | 81 | + | 
|  | 82 | +    XCTAssertEqual(serverContent.turnComplete, true) | 
|  | 83 | +    XCTAssertNil(serverContent.interrupted) | 
|  | 84 | +    XCTAssertNil(serverContent.generationComplete) | 
|  | 85 | + | 
|  | 86 | +    let modelTurn = try XCTUnwrap(serverContent.modelTurn) | 
|  | 87 | +    XCTAssertEqual(modelTurn.role, "model") | 
|  | 88 | +    XCTAssertEqual(modelTurn.parts.count, 1) | 
|  | 89 | +    let part = try XCTUnwrap(modelTurn.parts.first as? InlineDataPart) | 
|  | 90 | +    XCTAssertEqual(part.data, Data(repeating: 5, count: 5)) | 
|  | 91 | +    XCTAssertEqual(part.mimeType, "audio/pcm") | 
|  | 92 | + | 
|  | 93 | +    let metadata = try XCTUnwrap(serverContent.groundingMetadata) | 
|  | 94 | +    XCTAssertEqual(metadata.webSearchQueries, ["query1", "query2"]) | 
|  | 95 | +    XCTAssertEqual(metadata.groundingChunks.count, 1) | 
|  | 96 | +    let groundingChunk = try XCTUnwrap(metadata.groundingChunks.first) | 
|  | 97 | +    let webChunk = try XCTUnwrap(groundingChunk.web) | 
|  | 98 | +    XCTAssertEqual(webChunk.uri, "uri1") | 
|  | 99 | +    XCTAssertEqual(metadata.groundingSupports.count, 1) | 
|  | 100 | +    let groundingSupport = try XCTUnwrap(metadata.groundingSupports.first) | 
|  | 101 | +    XCTAssertEqual(groundingSupport.segment.startIndex, 0) | 
|  | 102 | +    XCTAssertEqual(groundingSupport.segment.partIndex, 0) | 
|  | 103 | +    XCTAssertEqual(groundingSupport.segment.endIndex, 10) | 
|  | 104 | +    XCTAssertEqual(groundingSupport.segment.text, "text") | 
|  | 105 | +    let searchEntryPoint = try XCTUnwrap(metadata.searchEntryPoint) | 
|  | 106 | +    XCTAssertEqual(searchEntryPoint.renderedContent, "html") | 
|  | 107 | + | 
|  | 108 | +    let inputTranscription = try XCTUnwrap(serverContent.inputTranscription) | 
|  | 109 | +    XCTAssertEqual(inputTranscription.text, "What day of the week is it?") | 
|  | 110 | + | 
|  | 111 | +    let outputTranscription = try XCTUnwrap(serverContent.outputTranscription) | 
|  | 112 | +    XCTAssertEqual(outputTranscription.text, "Today is friday") | 
|  | 113 | +  } | 
|  | 114 | + | 
|  | 115 | +  func testDecodeBidiGenerateContentServerMessage_toolCall() throws { | 
|  | 116 | +    let json = """ | 
|  | 117 | +    { | 
|  | 118 | +      "toolCall" : { | 
|  | 119 | +        "functionCalls" : [ | 
|  | 120 | +          { | 
|  | 121 | +            "name": "changeBackgroundColor", | 
|  | 122 | +            "id": "functionCall-12345-67890", | 
|  | 123 | +            "args" : { | 
|  | 124 | +              "color": "#F54927" | 
|  | 125 | +            } | 
|  | 126 | +          } | 
|  | 127 | +        ] | 
|  | 128 | +      } | 
|  | 129 | +    } | 
|  | 130 | +    """ | 
|  | 131 | +    let jsonData = try XCTUnwrap(json.data(using: .utf8)) | 
|  | 132 | + | 
|  | 133 | +    let serverMessage = try decoder.decode(BidiGenerateContentServerMessage.self, from: jsonData) | 
|  | 134 | +    guard case let .toolCall(toolCall) = serverMessage.messageType else { | 
|  | 135 | +      XCTFail("Decoded message is not a toolCall message.") | 
|  | 136 | +      return | 
|  | 137 | +    } | 
|  | 138 | + | 
|  | 139 | +    let functionCalls = try XCTUnwrap(toolCall.functionCalls) | 
|  | 140 | +    XCTAssertEqual(functionCalls.count, 1) | 
|  | 141 | +    let functionCall = try XCTUnwrap(functionCalls.first) | 
|  | 142 | +    XCTAssertEqual(functionCall.name, "changeBackgroundColor") | 
|  | 143 | +    XCTAssertEqual(functionCall.id, "functionCall-12345-67890") | 
|  | 144 | +    let args = try XCTUnwrap(functionCall.args) | 
|  | 145 | +    guard case let .string(color) = args["color"] else { | 
|  | 146 | +      XCTFail("Missing color argument") | 
|  | 147 | +      return | 
|  | 148 | +    } | 
|  | 149 | +    XCTAssertEqual(color, "#F54927") | 
|  | 150 | +  } | 
|  | 151 | + | 
|  | 152 | +  func testDecodeBidiGenerateContentServerMessage_toolCallCancellation() throws { | 
|  | 153 | +    let json = """ | 
|  | 154 | +    { | 
|  | 155 | +      "toolCallCancellation" : { | 
|  | 156 | +        "ids" : ["functionCall-12345-67890"] | 
|  | 157 | +      } | 
|  | 158 | +    } | 
|  | 159 | +    """ | 
|  | 160 | +    let jsonData = try XCTUnwrap(json.data(using: .utf8)) | 
|  | 161 | + | 
|  | 162 | +    let serverMessage = try decoder.decode(BidiGenerateContentServerMessage.self, from: jsonData) | 
|  | 163 | +    guard case let .toolCallCancellation(toolCallCancellation) = serverMessage.messageType else { | 
|  | 164 | +      XCTFail("Decoded message is not a toolCallCancellation message.") | 
|  | 165 | +      return | 
|  | 166 | +    } | 
|  | 167 | + | 
|  | 168 | +    let ids = try XCTUnwrap(toolCallCancellation.ids) | 
|  | 169 | +    XCTAssertEqual(ids, ["functionCall-12345-67890"]) | 
|  | 170 | +  } | 
|  | 171 | + | 
|  | 172 | +  func testDecodeBidiGenerateContentServerMessage_goAway() throws { | 
|  | 173 | +    let json = """ | 
|  | 174 | +    { | 
|  | 175 | +      "goAway" : { | 
|  | 176 | +        "timeLeft": "1.23456789s" | 
|  | 177 | +      } | 
|  | 178 | +    } | 
|  | 179 | +    """ | 
|  | 180 | +    let jsonData = try XCTUnwrap(json.data(using: .utf8)) | 
|  | 181 | + | 
|  | 182 | +    let serverMessage = try decoder.decode(BidiGenerateContentServerMessage.self, from: jsonData) | 
|  | 183 | +    guard case let .goAway(goAway) = serverMessage.messageType else { | 
|  | 184 | +      XCTFail("Decoded message is not a goAway message.") | 
|  | 185 | +      return | 
|  | 186 | +    } | 
|  | 187 | + | 
|  | 188 | +    XCTAssertEqual(goAway.timeLeft?.seconds, 1) | 
|  | 189 | +    XCTAssertEqual(goAway.timeLeft?.nanos, 234_567_890) | 
|  | 190 | +  } | 
|  | 191 | +} | 
0 commit comments