@Test
  public void sendMessageStreamWithText_throwsIllegalStateExceptionWhenFinishReasonIsNotSTOP()
      throws IOException {
    // (Arrange) Set up the responseStream
    ResponseStream<GenerateContentResponse> responseStream =
        new ResponseStream(new ResponseStreamIteratorWithHistory(mockServerStreamIterator));
    when(mockServerStreamIterator.hasNext()).thenReturn(true, true, false);
    when(mockServerStreamIterator.next())
        .thenReturn(
            RESPONSE_STREAM_CHUNK1_RESPONSE, RESPONSE_STREAM_CHUNK2_RESPONSE_WITHOUT_FINISH_REASON);

    // (Arrange) Set up the return value of the generateContentStream
    when(mockGenerativeModel.generateContentStream(
            Arrays.asList(ContentMaker.fromString(SAMPLE_MESSAGE1))))
        .thenReturn(responseStream);

    // (Act) send request, consume response
    ResponseStream<GenerateContentResponse> returndResponseStream =
        chat.sendMessageStream(SAMPLE_MESSAGE1);
    returndResponseStream.stream().forEach(element -> {});

    // (Act & Assert) get history and assert IllegalStateException for not having the right finish
    // reason.
    IllegalStateException thrown =
        assertThrows(IllegalStateException.class, () -> chat.getHistory());
    assertThat(thrown).hasMessageThat().isEqualTo("Rerun getHistory() to get cleaned history.");

    // Assert that the history can be fetched again and it's empty.
    List<Content> history = chat.getHistory();
    assertThat(history.size()).isEqualTo(0);
  }
Callable function getCurrentWeather for testing automatic function calling.
java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/generativeai/ChatSessionTest.java