chore(vertexai): Add corresponding unit test for the fork

packages/firebase_vertexai/firebase_vertexai/lib/src/vertex_api.dart

@@ -226,6 +226,36 @@ final class Candidate {
 
   /// Message for finish reason.
   final String? finishMessage;
+
+  /// The concatenation of the text parts of [content], if any.
+  ///
+  /// If this candidate was finished for a reason of [FinishReason.recitation]
+  /// or [FinishReason.safety], accessing this text will throw a
+  /// [GenerativeAIException].
+  ///
+  /// If [content] contains any text parts, this value is the concatenation of
+  /// the text.
+  ///
+  /// If [content] does not contain any text parts, this value is `null`.
+  String? get text {
+    if (finishReason case FinishReason.recitation || FinishReason.safety) {
+      final String suffix;
+      if (finishMessage case final message? when message.isNotEmpty) {
+        suffix = ': $message';
+      } else {
+        suffix = '';
+      }
+      throw VertexAIException(
+          'Candidate was blocked due to $finishReason$suffix');
+    }
+    return switch (content.parts) {
+      // Special case for a single TextPart to avoid iterable chain.
+      [TextPart(:final text)] => text,
+      final parts when parts.any((p) => p is TextPart) =>
+        parts.whereType<TextPart>().map((p) => p.text).join(),
+      _ => null,
+    };
+  }
 }
 
 /// Safety rating for a piece of content.
@@ -659,6 +689,8 @@ CountTokensResponse parseCountTokensResponse(Object jsonObject) {
         case {'totalBillableCharacters': final int totalBillableCharacters}) {
       return CountTokensResponse(totalTokens,
           totalBillableCharacters: totalBillableCharacters);
+    } else {
+      return CountTokensResponse(totalTokens);
     }
   }
   throw unhandledFormat('CountTokensResponse', jsonObject);

packages/firebase_vertexai/firebase_vertexai/lib/src/vertex_content.dart

@@ -14,6 +14,7 @@
 
 import 'dart:convert';
 import 'dart:typed_data';
+import 'vertex_error.dart';
 
 /// The base structured datatype containing multi-part content of a message.
 final class Content {
@@ -73,7 +74,7 @@ Content parseContent(Object jsonObject) {
           _ => null,
         },
         parts.map(_parsePart).toList()),
-    _ => throw FormatException('Unhandled Content format', jsonObject),
+    _ => throw unhandledFormat('Content', jsonObject),
   };
 }
 
@@ -100,7 +101,7 @@ Part _parsePart(Object? jsonObject) {
       throw UnimplementedError('FunctionResponse part not yet supported'),
     {'inlineData': {'mimeType': String _, 'data': String _}} =>
       throw UnimplementedError('inlineData content part not yet supported'),
-    _ => throw FormatException('Unhandled Part format', jsonObject),
+    _ => throw unhandledFormat('Part', jsonObject),
   };
 }
 

packages/firebase_vertexai/firebase_vertexai/lib/src/vertex_model.dart

@@ -90,6 +90,30 @@ final class GenerativeModel {
             httpClient: httpClient,
             requestHeaders: _firebaseTokens(appCheck, auth));
 
+  GenerativeModel._constructTestModel({
+    required String model,
+    required String location,
+    required FirebaseApp app,
+    FirebaseAppCheck? appCheck,
+    FirebaseAuth? auth,
+    List<SafetySetting>? safetySettings,
+    GenerationConfig? generationConfig,
+    List<Tool>? tools,
+    Content? systemInstruction,
+    ToolConfig? toolConfig,
+    ApiClient? apiClient,
+  })  : _model = _normalizeModelName(model),
+        _baseUri = _vertexUri(app, location),
+        _safetySettings = safetySettings ?? [],
+        _generationConfig = generationConfig,
+        _tools = tools,
+        _systemInstruction = systemInstruction,
+        _toolConfig = toolConfig,
+        _client = apiClient ??
+            HttpApiClient(
+                apiKey: app.options.apiKey,
+                requestHeaders: _firebaseTokens(appCheck, auth));
+
   final ({String prefix, String name}) _model;
   final List<SafetySetting> _safetySettings;
   final GenerationConfig? _generationConfig;
@@ -250,12 +274,8 @@ final class GenerativeModel {
   /// }
   /// ```
   Future<CountTokensResponse> countTokens(
-    Iterable<Content> contents, {
-    List<SafetySetting>? safetySettings,
-    GenerationConfig? generationConfig,
-    List<Tool>? tools,
-    ToolConfig? toolConfig,
-  }) async {
+    Iterable<Content> contents,
+  ) async {
     final parameters = <String, Object?>{
       'contents': contents.map((c) => c.toJson()).toList()
     };
@@ -337,3 +357,32 @@ GenerativeModel createGenerativeModel({
       tools: tools,
       toolConfig: toolConfig,
     );
+
+/// Creates a model with an overridden [ApiClient] for testing.
+///
+/// Package private test-only method.
+GenerativeModel createModelWithClient({
+  required FirebaseApp app,
+  required String location,
+  required String model,
+  required ApiClient client,
+  Content? systemInstruction,
+  FirebaseAppCheck? appCheck,
+  FirebaseAuth? auth,
+  GenerationConfig? generationConfig,
+  List<SafetySetting>? safetySettings,
+  List<Tool>? tools,
+  ToolConfig? toolConfig,
+}) =>
+    GenerativeModel._constructTestModel(
+        model: model,
+        app: app,
+        appCheck: appCheck,
+        auth: auth,
+        location: location,
+        safetySettings: safetySettings,
+        generationConfig: generationConfig,
+        systemInstruction: systemInstruction,
+        tools: tools,
+        toolConfig: toolConfig,
+        apiClient: client);

packages/firebase_vertexai/firebase_vertexai/pubspec.yaml

@@ -32,5 +32,6 @@ dev_dependencies:
   flutter_lints: ^4.0.0
   flutter_test:
     sdk: flutter
+  matcher: ^0.12.16
   mockito: ^5.0.0
   plugin_platform_interface: ^2.1.3

packages/firebase_vertexai/firebase_vertexai/test/utils/matchers.dart

@@ -0,0 +1,118 @@
+// Copyright 2024 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.
+import 'package:firebase_vertexai/firebase_vertexai.dart';
+import 'package:http/http.dart' as http;
+import 'package:matcher/matcher.dart';
+
+Matcher matchesPart(Part part) => switch (part) {
+      TextPart(text: final text) =>
+        isA<TextPart>().having((p) => p.text, 'text', text),
+      DataPart(mimeType: final mimeType, bytes: final bytes) => isA<DataPart>()
+          .having((p) => p.mimeType, 'mimeType', mimeType)
+          .having((p) => p.bytes, 'bytes', bytes),
+      FileData(mimeType: final mimeType, fileUri: final fileUri) =>
+        isA<FileData>()
+            .having((p) => p.mimeType, 'mimeType', mimeType)
+            .having((p) => p.fileUri, 'fileUri', fileUri),
+      FunctionCall(name: final name, args: final args) => isA<FunctionCall>()
+          .having((p) => p.name, 'name', name)
+          .having((p) => p.args, 'args', args),
+      FunctionResponse(name: final name, response: final response) =>
+        isA<FunctionResponse>()
+            .having((p) => p.name, 'name', name)
+            .having((p) => p.response, 'args', response),
+    };
+
+Matcher matchesContent(Content content) => isA<Content>()
+    .having((c) => c.role, 'role', content.role)
+    .having((c) => c.parts, 'parts', content.parts.map(matchesPart).toList());
+
+Matcher matchesCandidate(Candidate candidate) => isA<Candidate>().having(
+      (c) => c.content,
+      'content',
+      matchesContent(candidate.content),
+    );
+
+Matcher matchesGenerateContentResponse(GenerateContentResponse response) =>
+    isA<GenerateContentResponse>()
+        .having(
+          (r) => r.candidates,
+          'candidates',
+          response.candidates.map(matchesCandidate).toList(),
+        )
+        .having(
+          (r) => r.promptFeedback,
+          'promptFeedback',
+          response.promptFeedback == null
+              ? isNull
+              : matchesPromptFeedback(response.promptFeedback!),
+        );
+
+Matcher matchesPromptFeedback(
+  PromptFeedback promptFeedback,
+) =>
+    isA<PromptFeedback>()
+        .having((p) => p.blockReason, 'blockReason', promptFeedback.blockReason)
+        .having(
+          (p) => p.blockReasonMessage,
+          'blockReasonMessage',
+          promptFeedback.blockReasonMessage,
+        )
+        .having(
+          (p) => p.safetyRatings,
+          'safetyRatings',
+          unorderedMatches(
+              promptFeedback.safetyRatings.map(matchesSafetyRating)),
+        );
+
+Matcher matchesSafetyRating(SafetyRating safetyRating) => isA<SafetyRating>()
+    .having((s) => s.category, 'category', safetyRating.category)
+    .having((s) => s.probability, 'probability', safetyRating.probability);
+
+Matcher matchesEmbedding(ContentEmbedding embedding) =>
+    isA<ContentEmbedding>().having((e) => e.values, 'values', embedding.values);
+
+Matcher matchesEmbedContentResponse(EmbedContentResponse response) =>
+    isA<EmbedContentResponse>().having(
+      (r) => r.embedding,
+      'embedding',
+      matchesEmbedding(response.embedding),
+    );
+
+Matcher matchesBatchEmbedContentsResponse(
+  BatchEmbedContentsResponse response,
+) =>
+    isA<BatchEmbedContentsResponse>().having(
+      (r) => r.embeddings,
+      'embeddings',
+      response.embeddings.map(matchesEmbedding),
+    );
+
+Matcher matchesCountTokensResponse(CountTokensResponse response) =>
+    isA<CountTokensResponse>().having(
+      (r) => r.totalTokens,
+      'totalTokens',
+      response.totalTokens,
+    );
+
+Matcher matchesRequest(http.Request request) => isA<http.Request>()
+    .having((r) => r.headers, 'headers', request.headers)
+    .having((r) => r.method, 'method', request.method)
+    .having((r) => r.bodyBytes, 'bodyBytes', request.bodyBytes)
+    .having((r) => r.url, 'url', request.url);
+
+Matcher matchesBaseRequest(http.BaseRequest request) => isA<http.BaseRequest>()
+    .having((r) => r.headers, 'headers', request.headers)
+    .having((r) => r.method, 'method', request.method)
+    .having((r) => r.url, 'url', request.url);

packages/firebase_vertexai/firebase_vertexai/test/utils/stub_client.dart

@@ -0,0 +1,93 @@
+// Copyright 2024 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.
+
+import 'dart:collection';
+
+import 'package:firebase_vertexai/src/vertex_client.dart';
+
+class ClientController {
+  final _client = _ControlledClient();
+  ApiClient get client => _client;
+
+  /// Run [body] and return [response] for a single call to
+  /// [ApiClient.streamRequest].
+  ///
+  /// Check expectations for the request URI and JSON payload with the
+  /// [verifyRequest] callback.
+  Future<T> checkRequest<T>(
+    Future<T> Function() body, {
+    required Map<String, Object?> response,
+    void Function(Uri, Map<String, Object?>)? verifyRequest,
+  }) async {
+    _client._requestExpectations.addLast(verifyRequest);
+    _client._responses.addLast([response]);
+    final result = await body();
+    assert(_client._responses.isEmpty);
+    return result;
+  }
+
+  /// Run [body] and return [responses] for a single call to
+  /// [ApiClient.streamRequest].
+  ///
+  /// Check expectations for the request URI and JSON payload with the
+  /// [verifyRequest] callback.
+  Future<T> checkStreamRequest<T>(
+    Future<T> Function() body, {
+    required Iterable<Map<String, Object?>> responses,
+    void Function(Uri, Map<String, Object?>)? verifyRequest,
+  }) async {
+    _client._requestExpectations.addLast(verifyRequest);
+    _client._responses.addLast(responses.toList());
+    final result = await body();
+    assert(_client._responses.isEmpty);
+    return result;
+  }
+}
+
+final class _ControlledClient implements ApiClient {
+  final _requestExpectations =
+      Queue<void Function(Uri, Map<String, Object?>)?>();
+  final _responses = Queue<List<Map<String, Object?>>>();
+
+  @override
+  Future<Map<String, Object?>> makeRequest(
+    Uri uri,
+    Map<String, Object?> body,
+  ) async {
+    _requestExpectations.removeFirst()?.call(uri, body);
+    return _responses.removeFirst().single;
+  }
+
+  @override
+  Stream<Map<String, Object?>> streamRequest(
+    Uri uri,
+    Map<String, Object?> body,
+  ) {
+    _requestExpectations.removeFirst()?.call(uri, body);
+    return Stream.fromIterable(_responses.removeFirst());
+  }
+}
+
+const Map<String, Object?> arbitraryGenerateContentResponse = {
+  'candidates': [
+    {
+      'content': {
+        'role': 'model',
+        'parts': [
+          {'text': 'Some Response'},
+        ],
+      },
+    },
+  ],
+};