firebase-ios-sdk/GeminiLanguageModel/Tests/GeminiLanguageModelTests/IntegrationTests/BasicContentGenerationIntegrationTests.swift

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#if canImport(FoundationModels) && compiler(>=6.4)
  import Foundation
  import FoundationModels
  import GeminiAPIClient
  import GeminiTestUtilities
  import Testing

  @testable import GeminiLanguageModel

  /// Integration tests for content generation using `GeminiLanguageModel` with a basic feature set.
  @Suite(
    "Basic Content Generation Integration Tests",
    .requireFoundationModels,
    .tags(.integration)
  )
  struct BasicContentGenerationIntegrationTests {
    @Test(
      .requireIntegrationTestingBackend,
      arguments: IntegrationTestingBackend.availableBackends
    )
    @available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
    func sessionRespond(backend: IntegrationTestingBackend) async throws {
      let model = try await backend.makeModel()
      let session = LanguageModelSession(model: model)

      let response = try await session.respond(to: "Reply with the single word 'HELLO'.")

      #expect(!response.content.isEmpty)
      #expect(response.content.localizedCaseInsensitiveContains("HELLO"))
      #expect(response.usage.totalTokenCount > 0)
      #expect(response.usage.input.totalTokenCount > 0)
      #expect(response.usage.output.totalTokenCount > 0)
    }

    @Test(
      .requireIntegrationTestingBackend,
      arguments: IntegrationTestingBackend.availableBackends
    )
    @available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
    func sessionRespondMultiTurn(backend: IntegrationTestingBackend) async throws {
      let model = try await backend.makeModel()
      let session = LanguageModelSession(model: model)

      _ = try await session.respond(to: "My favorite color is teal.")
      let response = try await session.respond(
        to: "What is my favorite color? Answer in one word."
      )

      #expect(!response.content.isEmpty)
      #expect(response.content.localizedCaseInsensitiveContains("teal"))
      #expect(response.usage.totalTokenCount > 0)
      #expect(response.usage.input.totalTokenCount > 0)
      #expect(response.usage.output.totalTokenCount > 0)
    }

    @Test(
      .requireIntegrationTestingBackend,
      arguments: IntegrationTestingBackend.availableBackends
    )
    @available(macOS 27.0, iOS 27.0, watchOS 27.0, visionOS 27.0, *)
    func sessionStreamResponse(backend: IntegrationTestingBackend) async throws {
      let model = try await backend.makeModel()
      let session = LanguageModelSession(model: model)

      let stream = session.streamResponse(
        to: "List the numbers 1 through 5, one per line."
      )
      var finalContent = ""
      var snapshotCount = 0
      var finalUsage: LanguageModelSession.Usage?
      for try await snapshot in stream {
        finalContent = snapshot.content
        snapshotCount += 1
        finalUsage = snapshot.usage
      }

      #expect(snapshotCount > 1)
      #expect(!finalContent.isEmpty)
      #expect(finalContent.contains("1"))
      #expect(finalContent.contains("5"))
      let usage = try #require(finalUsage)
      #expect(usage.totalTokenCount > 0)
      #expect(usage.input.totalTokenCount > 0)
      #expect(usage.output.totalTokenCount > 0)
    }
  }
#endif  // canImport(FoundationModels) && compiler(>=6.4)