1+ // Copyright 2024 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 Foundation
16+ import XCTest
17+
18+ @testable import FirebaseAILogic
19+
20+ @available ( iOS 15 . 0 , macOS 12 . 0 , macCatalyst 15 . 0 , tvOS 15 . 0 , watchOS 8 . 0 , * )
21+ final class ModelContentTests : XCTestCase {
22+ // MARK: - ModelContent Initialization with Part Types
23+
24+ func testInitWithExecutableCodePart( ) {
25+ let executableCodePart = ExecutableCodePart ( language: . python, code: " print('hello') " )
26+
27+ let content = ModelContent ( role: " model " , parts: [ executableCodePart] )
28+
29+ XCTAssertEqual ( content. role, " model " )
30+ XCTAssertEqual ( content. parts. count, 1 )
31+ guard let resultPart = content. parts. first as? ExecutableCodePart else {
32+ XCTFail ( " Expected ExecutableCodePart " )
33+ return
34+ }
35+ XCTAssertEqual ( resultPart. language, . python)
36+ XCTAssertEqual ( resultPart. code, " print('hello') " )
37+ }
38+
39+ func testInitWithCodeExecutionResultPart( ) {
40+ let codeExecutionResultPart = CodeExecutionResultPart ( outcome: . ok, output: " hello " )
41+
42+ let content = ModelContent ( role: " model " , parts: [ codeExecutionResultPart] )
43+
44+ XCTAssertEqual ( content. role, " model " )
45+ XCTAssertEqual ( content. parts. count, 1 )
46+ guard let resultPart = content. parts. first as? CodeExecutionResultPart else {
47+ XCTFail ( " Expected CodeExecutionResultPart " )
48+ return
49+ }
50+ XCTAssertEqual ( resultPart. outcome, . ok)
51+ XCTAssertEqual ( resultPart. output, " hello " )
52+ }
53+
54+ func testInitWithMixedPartsIncludingExecutableCode( ) {
55+ let textPart = TextPart ( " Here is some code: " )
56+ let executableCodePart = ExecutableCodePart ( language: . python, code: " x = 1 + 2 " )
57+ let codeExecutionResultPart = CodeExecutionResultPart ( outcome: . ok, output: " 3 " )
58+
59+ let content = ModelContent (
60+ role: " model " ,
61+ parts: [ textPart, executableCodePart, codeExecutionResultPart]
62+ )
63+
64+ XCTAssertEqual ( content. role, " model " )
65+ XCTAssertEqual ( content. parts. count, 3 )
66+
67+ // Verify each part type
68+ XCTAssertTrue ( content. parts [ 0 ] is TextPart )
69+ XCTAssertTrue ( content. parts [ 1 ] is ExecutableCodePart )
70+ XCTAssertTrue ( content. parts [ 2 ] is CodeExecutionResultPart )
71+ }
72+
73+ func testInitWithExecutableCodePartPreservesThoughtMetadata( ) {
74+ // Test that thought-related metadata is preserved through the conversion
75+ let executableCodePart = ExecutableCodePart ( language: . python, code: " print('test') " )
76+
77+ let content = ModelContent ( role: " model " , parts: [ executableCodePart] )
78+
79+ guard let resultPart = content. parts. first as? ExecutableCodePart else {
80+ XCTFail ( " Expected ExecutableCodePart " )
81+ return
82+ }
83+ // Verify the part maintains its properties after round-trip
84+ XCTAssertEqual ( resultPart. language, executableCodePart. language)
85+ XCTAssertEqual ( resultPart. code, executableCodePart. code)
86+ }
87+
88+ func testInitWithCodeExecutionResultPartWithDeadlockedOutcome( ) {
89+ let codeExecutionResultPart = CodeExecutionResultPart ( outcome: . deadlineExceeded, output: nil )
90+
91+ let content = ModelContent ( role: " model " , parts: [ codeExecutionResultPart] )
92+
93+ guard let resultPart = content. parts. first as? CodeExecutionResultPart else {
94+ XCTFail ( " Expected CodeExecutionResultPart " )
95+ return
96+ }
97+ XCTAssertEqual ( resultPart. outcome, . deadlineExceeded)
98+ XCTAssertNil ( resultPart. output)
99+ }
100+ }
0 commit comments