content: /** Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.atlas.impala;
  
  import org.apache.atlas.impala.hook.ImpalaLineageHook;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.testng.annotations.AfterMethod;
  import org.testng.annotations.BeforeMethod;
  import org.testng.annotations.Test;
  import org.mockito.Mock;
  import org.mockito.MockitoAnnotations;
  
  import java.io.ByteArrayOutputStream;
  import java.io.File;
  import java.io.FileWriter;
  import java.io.IOException;
  import java.io.PrintStream;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.util.Comparator;
  
  import static org.mockito.ArgumentMatchers.anyString;
  import static org.mockito.Mockito.*;
          import static org.testng.Assert.*;
  
  public class ImpalaLineageToolTest {
  
      private static final Logger LOG = LoggerFactory.getLogger(ImpalaLineageToolTest.class);
  
      @Mock
      private ImpalaLineageHook mockImpalaLineageHook;
  
      private Path tempDir;
      private final PrintStream originalOut = System.out;
      private ByteArrayOutputStream testOut;
  
      @BeforeMethod
      public void setUp() throws IOException {
          MockitoAnnotations.initMocks(this);
  
          tempDir = Files.createTempDirectory("impala-lineage-test");
  
          testOut = new ByteArrayOutputStream();
          System.setOut(new PrintStream(testOut));
      }
  
      @AfterMethod
      public void tearDown() throws IOException {
  
          if (tempDir != null && Files.exists(tempDir)) {
              Files.walk(tempDir)
                      .sorted(Comparator.reverseOrder())
                      .map(Path::toFile)
                      .forEach(File::delete);
          }
  
          System.setOut(originalOut);
      }
  
      @Test
      public void testConstructorWithValidArguments() {
          String[] args = {"-d", "/test/directory", "-p", "test_prefix"};
  
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          assertNotNull(tool);
      }
  
      @Test(expectedExceptions = RuntimeException.class)
      public void testConstructorWithInvalidArguments() {
          String[] args = {"-invalid", "argument"};
  
          new ImpalaLineageTool(args);
      }
  
      @Test
      public void testGetCurrentFilesWithNoFiles() throws IOException {
          String[] args = {"-d", tempDir.toString(), "-p", "nonexistent"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          File[] files = tool.getCurrentFiles();
  
          assertNotNull(files);
          assertEquals(files.length, 0);
      }
  
      @Test
      public void testGetCurrentFilesWithSingleFile() throws IOException {
  
          File testFile = new File(tempDir.toFile(), "test_lineage.log");
          testFile.createNewFile();
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          File[] files = tool.getCurrentFiles();
  
          assertNotNull(files);
          assertEquals(files.length, 1);
          assertEquals(files[0].getName(), "test_lineage.log");
      }
  
      @Test
      public void testGetCurrentFilesWithMultipleFilesSorted() throws IOException {
  
          File file1 = new File(tempDir.toFile(), "test_file1.log");
          File file2 = new File(tempDir.toFile(), "test_file2.log");
          File file3 = new File(tempDir.toFile(), "test_file3.log");
  
          file1.createNewFile();
          file2.createNewFile();
          file3.createNewFile();
  
  
          long baseTime = System.currentTimeMillis();
          file1.setLastModified(baseTime - 3000); // 3 seconds ago
          file2.setLastModified(baseTime - 2000); // 2 seconds ago
          file3.setLastModified(baseTime - 1000); // 1 second ago
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          File[] files = tool.getCurrentFiles();
  
          assertNotNull(files);
          assertEquals(files.length, 3);
  
          assertTrue(files[0].lastModified() <= files[1].lastModified());
          assertTrue(files[1].lastModified() <= files[2].lastModified());
      }
  
      @Test
      public void testGetCurrentFilesWithNonExistentDirectory() {
          String[] args = {"-d", "/nonexistent/directory", "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          File[] files = tool.getCurrentFiles();
  
          assertNotNull(files);
          assertEquals(files.length, 0);
      }
  
      @Test
      public void testDeleteLineageAndWalSuccess() throws IOException {
          File lineageFile = new File(tempDir.toFile(), "test_lineage.log");
          File walFile = new File(tempDir.toFile(), "test_wal.wal");
  
          lineageFile.createNewFile();
          walFile.createNewFile();
  
          assertTrue(lineageFile.exists());
          assertTrue(walFile.exists());
  
          ImpalaLineageTool.deleteLineageAndWal(lineageFile, walFile.getAbsolutePath());
  
          assertFalse(lineageFile.exists());
          assertFalse(walFile.exists());
      }
  
      @Test
      public void testDeleteLineageAndWalNonExistentFiles() {
          File nonExistentFile = new File(tempDir.toFile(), "nonexistent.log");
          String nonExistentWalPath = tempDir.resolve("nonexistent.wal").toString();
  
          ImpalaLineageTool.deleteLineageAndWal(nonExistentFile, nonExistentWalPath);
      }
  
      @Test
      public void testImportHImpalaEntitiesWithNewWalFile() throws Exception {
  
          File lineageFile = new File(tempDir.toFile(), "test_lineage.log");
          String testContent = "test lineage content";
  
          try (FileWriter writer = new FileWriter(lineageFile)) {
              writer.write(testContent);
          }
  
          String walPath = tempDir.resolve("test.wal").toString();
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
  
          doNothing().when(mockImpalaLineageHook).process(anyString());
  
          tool.importHImpalaEntities(mockImpalaLineageHook, lineageFile.getAbsolutePath(), walPath);
  
          verify(mockImpalaLineageHook, times(1)).process(testContent);
  
          File walFile = new File(walPath);
          assertTrue(walFile.exists());
      }
  
      @Test
      public void testImportHImpalaEntitiesWithExistingWalFile() throws Exception {
  
          File lineageFile = new File(tempDir.toFile(), "test_lineage.log");
          String testContent = "test lineage content";
  
          try (FileWriter writer = new FileWriter(lineageFile)) {
              writer.write(testContent);
          }
  
  
          String walPath = tempDir.resolve("test.wal").toString();
          File walFile = new File(walPath);
          try (FileWriter writer = new FileWriter(walFile)) {
              writer.write("0, "+  lineageFile.getAbsolutePath());
          }
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          doNothing().when(mockImpalaLineageHook).process(anyString());
  
          tool.importHImpalaEntities(mockImpalaLineageHook, lineageFile.getAbsolutePath(), walPath);
  
          verify(mockImpalaLineageHook, times(1)).process(testContent);
      }
  
      @Test
      public void testImportHImpalaEntitiesWithProcessingFailure() throws Exception {
  
          File lineageFile = new File(tempDir.toFile(), "test_lineage.log");
          String testContent = "test lineage content";
  
          try (FileWriter writer = new FileWriter(lineageFile)) {
              writer.write(testContent);
          }
  
          String walPath = tempDir.resolve("test.wal").toString();
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          doThrow(new RuntimeException("Processing failed")).when(mockImpalaLineageHook).process(anyString());
  
          tool.importHImpalaEntities(mockImpalaLineageHook, lineageFile.getAbsolutePath(), walPath);
  
          verify(mockImpalaLineageHook, times(1)).process(testContent);
  
          File walFile = new File(walPath);
          assertTrue(walFile.exists());
          String walContent = new String(Files.readAllBytes(walFile.toPath()));
          assertEquals(walContent.trim(), "0, "+  lineageFile.getAbsolutePath());
      }
  
  
      @Test
      public void testMainWithIncorrectNumberOfArguments() {
          String[] args = {"-d", tempDir.toString()};
  
          ImpalaLineageTool.main(args);
  
          String output = testOut.toString();
          assertTrue(output.contains("wrong number of arguments"));
          assertTrue(output.contains("Usage: import-impala.sh"));
      }
  
  
      @Test
      public void testRunWithMultipleFiles() throws IOException {
          // Create multiple test files
          File file1 = new File(tempDir.toFile(), "test_file1.log");
          File file2 = new File(tempDir.toFile(), "test_file2.log");
          File file3 = new File(tempDir.toFile(), "test_file3.log");
  
          file1.createNewFile();
          try (FileWriter writer = new FileWriter(file1)) {
              writer.write("content1");
          }
  
          file2.createNewFile();
          try (FileWriter writer = new FileWriter(file2)) {
              writer.write("content2");
          }
  
          file3.createNewFile();
          try (FileWriter writer = new FileWriter(file3)) {
              writer.write("content3");
          }
  
          long baseTime = System.currentTimeMillis();
          file1.setLastModified(baseTime - 3000);
          file2.setLastModified(baseTime - 2000);
          file3.setLastModified(baseTime - 1000);
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          try {
              tool.run();
          } catch (Exception e) {
              LOG.info("Error running test case ImpalaLineageToolTest.testRunWithMultipleFiles()");
          }
  
          assertFalse(file1.exists());
          assertFalse(file2.exists());
          assertTrue(file3.exists());
      }
  
      @Test
      public void testRunWithNoFiles() {
          String[] args = {"-d", tempDir.toString(), "-p", "nonexistent"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          tool.run();
      }
  
      @Test
      public void testRunWithSingleFile() throws IOException {
          File testFile = new File(tempDir.toFile(), "test_single.log");
          testFile.createNewFile();
          try (FileWriter writer = new FileWriter(testFile)) {
              writer.write("single file content");
          }
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          try {
              tool.run();
          } catch (Exception e) {
              LOG.info("Error running test case ImpalaLineageToolTest.testRunWithSingleFile()");
          }
  
          assertTrue(testFile.exists());
      }
  
      @Test
      public void testProcessImpalaLineageHookSuccess() throws Exception {
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          java.lang.reflect.Method method = ImpalaLineageTool.class.getDeclaredMethod(
                  "processImpalaLineageHook",
                  ImpalaLineageHook.class,
                  java.util.List.class
          );
          method.setAccessible(true);
  
          doNothing().when(mockImpalaLineageHook).process(anyString());
  
          java.util.List<String> lineageList = java.util.Arrays.asList("query1", "query2", "query3");
          boolean result = (Boolean) method.invoke(tool, mockImpalaLineageHook, lineageList);
  
          assertTrue(result);
          verify(mockImpalaLineageHook, times(3)).process(anyString());
      }
  
      @Test
      public void testProcessImpalaLineageHookWithFailures() throws Exception {
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          // Use reflection to access the private method
          java.lang.reflect.Method method = ImpalaLineageTool.class.getDeclaredMethod(
                  "processImpalaLineageHook",
                  ImpalaLineageHook.class,
                  java.util.List.class
          );
          method.setAccessible(true);
  
          doNothing().when(mockImpalaLineageHook).process("query1");
          doThrow(new RuntimeException("Processing failed")).when(mockImpalaLineageHook).process("query2");
          doNothing().when(mockImpalaLineageHook).process("query3");
  
          java.util.List<String> lineageList = java.util.Arrays.asList("query1", "query2", "query3");
          boolean result = (Boolean) method.invoke(tool, mockImpalaLineageHook, lineageList);
  
          assertFalse(result);
          verify(mockImpalaLineageHook, times(3)).process(anyString());
      }
  
      @Test
      public void testProcessImpalaLineageHookWithEmptyList() throws Exception {
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          java.lang.reflect.Method method = ImpalaLineageTool.class.getDeclaredMethod(
                  "processImpalaLineageHook",
                  ImpalaLineageHook.class,
                  java.util.List.class
          );
          method.setAccessible(true);
  
          java.util.List<String> emptyList = java.util.Collections.emptyList();
          boolean result = (Boolean) method.invoke(tool, mockImpalaLineageHook, emptyList);
  
          assertTrue(result);
          verify(mockImpalaLineageHook, never()).process(anyString());
      }
  
      @Test
      public void testGetCurrentFilesWithMultipleFilesSortedByModificationTime() throws IOException, InterruptedException {
          // Create multiple test files and ensure different modification times
          File file1 = new File(tempDir.toFile(), "test_old.log");
          file1.createNewFile();
          file1.setLastModified(System.currentTimeMillis() - 3000); // 3 seconds ago
  
          File file2 = new File(tempDir.toFile(), "test_medium.log");
          file2.createNewFile();
          file2.setLastModified(System.currentTimeMillis() - 2000); // 2 seconds ago
  
          File file3 = new File(tempDir.toFile(), "test_new.log");
          file3.createNewFile();
          file3.setLastModified(System.currentTimeMillis() - 1000); // 1 second ago
  
          String[] args = {"-d", tempDir.toString(), "-p", "test"};
          ImpalaLineageTool tool = new ImpalaLineageTool(args);
  
          File[] files = tool.getCurrentFiles();
  
          assertNotNull(files);
          assertEquals(files.length, 3);
          assertEquals(files[0].getName(), "test_old.log");
          assertEquals(files[1].getName(), "test_medium.log");
          assertEquals(files[2].getName(), "test_new.log");
      }
  
  }
  