2121import static com .google .common .truth .Truth .assertThat ;
2222import static com .google .common .truth .Truth .assertWithMessage ;
2323
24- import java .io .ByteArrayOutputStream ;
2524import java .io .IOException ;
26- import java .io .PrintStream ;
2725import org .apache .http .HttpException ;
28- import org .junit .Before ;
26+ import org .junit .BeforeClass ;
2927import org .junit .Test ;
3028import org .junit .runner .RunWith ;
3129import org .junit .runners .JUnit4 ;
3230
3331@ RunWith (JUnit4 .class )
3432public class TextGenerationIT {
3533
36- private ByteArrayOutputStream bout ;
37- private PrintStream originalOut ;
38-
34+ private String MODEL_ID ="gemini-2.0-flash-001" ;
3935
4036 // Check if the required environment variables are set.
4137 public static void requireEnvVar (String envVarName ) {
@@ -44,21 +40,18 @@ public static void requireEnvVar(String envVarName) {
4440 .isNotEmpty ();
4541 }
4642
47- @ Before
48- public void setUp () {
49- originalOut = System .out ;
50- bout = new ByteArrayOutputStream ();
51- System .setOut (new PrintStream (bout ));
43+ @ BeforeClass
44+ public static void setUp () {
45+ requireEnvVar ("GOOGLE_APPLICATION_CREDENTIALS" );
46+ requireEnvVar ("GOOGLE_CLOUD_PROJECT" );
47+ requireEnvVar ("GOOGLE_CLOUD_LOCATION" );
48+ requireEnvVar ("GOOGLE_GENAI_USE_VERTEXAI" );
5249 }
5350
54- public void tearDown () {
55- System .setOut (originalOut );
56- }
57-
5851 @ Test
59- public void testTextgenWithTxt () throws IOException , HttpException {
60- TxtgenWithTxt . main ( new String []{}) ;
61- tearDown ( );
62- assertThat (bout . toString () ).isNotEmpty ();
52+ public void testTextGeneration () throws IOException , HttpException {
53+ String prompt = "How does AI work?" ;
54+ String response = TextGeneration . generateContent ( MODEL_ID , prompt );
55+ assertThat (response ).isNotEmpty ();
6356 }
6457}
0 commit comments