Skip to content

Commit 6eba765

Browse files
committed
test(FirebaseAI): Add unit tests for ExecutableCodePart and CodeExecutionResultPart in ModelContent
Add tests to verify ModelContent can be initialized with: - ExecutableCodePart - CodeExecutionResultPart - Mixed parts including these types These tests ensure the bug fix is protected against regression.
1 parent 9a731da commit 6eba765

File tree

1 file changed

+100
-0
lines changed

1 file changed

+100
-0
lines changed
Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
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

Comments
 (0)