/*
 * Copyright 2020 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.
 */

package com.example.automl;

import static com.google.common.truth.Truth.assertThat;
import static junit.framework.TestCase.assertNotNull;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import org.junit.After;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

@RunWith(JUnit4.class)
public class GetOperationStatusTest {
  private static final String PROJECT_ID = System.getenv("AUTOML_PROJECT_ID");
  private String operationId;
  private ByteArrayOutputStream bout;
  private PrintStream out;

  private static void requireEnvVar(String varName) {
    assertNotNull(
        System.getenv(varName),
        "Environment variable '%s' is required to perform these tests.".format(varName));
  }

  @BeforeClass
  public static void checkRequirements() {
    requireEnvVar("GOOGLE_APPLICATION_CREDENTIALS");
    requireEnvVar("AUTOML_PROJECT_ID");
  }

  @Before
  public void setUp() throws IOException {
    bout = new ByteArrayOutputStream();
    out = new PrintStream(bout);
    System.setOut(out);

    ListOperationStatus.listOperationStatus(PROJECT_ID);
    String got = bout.toString();
    operationId = got.split("\n")[1].split(":")[1].trim();
    assertThat(got).contains("Operation details:");

    bout = new ByteArrayOutputStream();
    out = new PrintStream(bout);
    System.setOut(out);
  }

  @After
  public void tearDown() {
    System.setOut(null);
  }

  @Test
  public void testGetOperationStatus() throws IOException {
    GetOperationStatus.getOperationStatus(operationId);
    String got = bout.toString();
    assertThat(got).contains("Operation details:");
  }
}

@@ -735,93 +735,533 @@
 ort 
-java.io.ByteArrayOutputStream;%0Aimport java.io.IOException;%0Aimport java.io.PrintStream
+com.google.api.gax.longrunning.OperationFuture;%0Aimport com.google.cloud.automl.v1.AutoMlClient;%0Aimport com.google.cloud.automl.v1.DatasetName;%0Aimport com.google.cloud.automl.v1.GcsSource;%0Aimport com.google.cloud.automl.v1.InputConfig;%0Aimport com.google.cloud.automl.v1.OperationMetadata;%0Aimport com.google.protobuf.Empty;%0Aimport java.io.ByteArrayOutputStream;%0Aimport java.io.IOException;%0Aimport java.io.PrintStream;%0Aimport java.util.Arrays;%0Aimport java.util.concurrent.ExecutionException;%0Aimport java.util.concurrent.TimeUnit
 ;%0Aim
@@ -1567,16 +1567,173 @@
 T_ID%22);%0A
+  private static final String DATASET_ID = %22TRL5394674636845744128%22;%0A  private static final String BUCKET = %22gs://translate_data_exported/translation.csv%22;%0A%0A
   privat
@@ -1750,18 +1750,24 @@
 peration
-Id
+FullName
 ;%0A  priv
@@ -2228,24 +2228,66 @@
  IOException
+, ExecutionException, InterruptedException
  %7B%0A    bout 
@@ -2383,205 +2383,692 @@
     
-ListOperationStatus.listOperationStatus(PROJEC
+// start a export data into dataset and cancel it before it finishes.%0A    try (AutoMlClient client = AutoMlClient.create()) %7B%0A      // Get the complete path of the dataset.%0A      DatasetName datasetFullId = DatasetName.of(PROJECT_ID, %22us-central1%22, DATASE
 T_ID);%0A
+%0A 
     
-String got = bout.toString
+ GcsSource gcsSource =%0A          GcsSource.newBuilder().addAllInputUris(Arrays.asList(BUCKET.split(%22,%22))).build
 ();%0A
+%0A 
     
-operationId = got.split(%22%5Cn%22)%5B1%5D.split(%22:%22)%5B1%5D.trim
+ InputConfig inputConfig = InputConfig.newBuilder().setGcsSource(gcsSource).build
 ();%0A
+%0A 
     
-assertThat(got).contains(%22Operation details:%22);
+ // Start the import LRO job%0A      OperationFuture%3CEmpty, OperationMetadata%3E operation =%0A          client.importDataAsync(datasetFullId, inputConfig);%0A%0A      operationFullName = operation.getName();%0A    %7D
 %0A%0A  
@@ -3171,74 +3171,8 @@
  %7D%0A%0A
-  @After%0A  public void tearDown() %7B%0A    System.setOut(null);%0A  %7D%0A%0A
   @T
@@ -3286,18 +3286,24 @@
 peration
-Id
+FullName
 );%0A    S
@@ -3386,11 +3386,362 @@
 :%22);%0A  %7D
+%0A%0A  @After%0A  public void tearDown() throws IOException, InterruptedException %7B%0A    try (AutoMlClient client = AutoMlClient.create()) %7B%0A      // terminate export data LRO.%0A      client.getOperationsClient().cancelOperation(operationFullName);%0A      client.getOperationsClient().awaitTermination(5, TimeUnit.SECONDS);%0A    %7D%0A%0A    System.setOut(null);%0A  %7D
 %0A%7D%0A

refactor automl test (#3588)