diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
index aa4123d5a9e..d25b6b60d30 100644
--- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
+++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift
@@ -222,7 +222,9 @@ struct GenerateContentIntegrationTests {
 
     let candidate = try #require(response.candidates.first)
     let thoughtParts = candidate.content.parts.compactMap { $0.isThought ? $0 : nil }
-    #expect(thoughtParts.isEmpty != (thinkingConfig.includeThoughts ?? false))
+    let thoughtSignature = candidate.content.internalParts.first?.thoughtSignature
+    #expect(thoughtSignature != nil || thoughtParts
+      .isEmpty != (thinkingConfig.includeThoughts ?? false))
 
     let usageMetadata = try #require(response.usageMetadata)
     #expect(usageMetadata.promptTokenCount.isEqual(to: 13, accuracy: tokenCountAccuracy))