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.hook.events;
  
  import org.apache.atlas.impala.hook.AtlasImpalaHookContext;
  import org.apache.atlas.impala.model.ImpalaDataType;
  import org.apache.atlas.impala.model.ImpalaNode;
  import org.apache.atlas.impala.model.ImpalaOperationType;
  import org.apache.atlas.impala.model.ImpalaQuery;
  import org.apache.atlas.impala.model.ImpalaVertexType;
  import org.apache.atlas.impala.model.LineageVertex;
  import org.apache.atlas.impala.model.LineageVertexMetadata;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
  import org.apache.atlas.model.instance.AtlasObjectId;
  import org.apache.atlas.model.notification.HookNotification;
  import org.mockito.Mock;
  import org.mockito.MockitoAnnotations;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.testng.annotations.BeforeMethod;
  import org.testng.annotations.DataProvider;
  import org.testng.annotations.Test;
  
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.List;
  
  import static org.apache.atlas.impala.hook.events.BaseImpalaEvent.*;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.ArgumentMatchers.anyString;
  import static org.mockito.Mockito.when;
  import static org.testng.Assert.*;
  
  public class BaseImpalaEventTest {
      private static final Logger LOG = LoggerFactory.getLogger(BaseImpalaEventTest.class);
  
      @Mock
      private AtlasImpalaHookContext mockContext;
  
      @Mock
      private ImpalaQuery mockQuery;
  
      private static final String CLUSTER_NAME = "test_cluster";
      private static final String DB_NAME = "test_database";
      private static final String TABLE_NAME = "test_table";
      private static final String COLUMN_NAME = "test_column";
      private static final String USER_NAME = "test_user";
      private static final String HOST_NAME = "test_host";
      private static final String QUERY_TEXT = "SELECT * FROM test_table";
      private static final String QUERY_ID = "test_query_123";
      private static final long TIMESTAMP = 1554750072L;
      private static final long END_TIME = 1554750554L;
  
      private static class TestBaseImpalaEvent extends BaseImpalaEvent {
          public TestBaseImpalaEvent(AtlasImpalaHookContext context) {
              super(context);
          }
  
          @Override
          public List<HookNotification> getNotificationMessages() throws Exception {
              return Collections.emptyList();
          }
      }
  
      private TestBaseImpalaEvent baseImpalaEvent;
  
      @BeforeMethod
      public void setUp() throws Exception {
          MockitoAnnotations.initMocks(this);
  
          when(mockContext.getUserName()).thenReturn(USER_NAME);
          when(mockContext.getHostName()).thenReturn(HOST_NAME);
          when(mockContext.getMetadataNamespace()).thenReturn(CLUSTER_NAME);
          when(mockContext.getQueryStr()).thenReturn(QUERY_TEXT);
          when(mockContext.getImpalaOperationType()).thenReturn(ImpalaOperationType.QUERY);
          when(mockContext.getLineageQuery()).thenReturn(mockQuery);
  
          when(mockQuery.getQueryId()).thenReturn(QUERY_ID);
          when(mockQuery.getTimestamp()).thenReturn(TIMESTAMP);
          when(mockQuery.getEndTime()).thenReturn(END_TIME);
  
          setupQualifiedNameMocks();
  
          baseImpalaEvent = new TestBaseImpalaEvent(mockContext);
      }
  
      private void setupQualifiedNameMocks() {
          when(mockContext.getQualifiedNameForDb(anyString())).thenAnswer(invocation -> {
              String dbName = invocation.getArgument(0);
              return dbName.toLowerCase()  +"@"+  CLUSTER_NAME;
          });
  
          when(mockContext.getQualifiedNameForTable(anyString())).thenAnswer(invocation -> {
              String tableQualifiedName = invocation.getArgument(0);
              return tableQualifiedName.toLowerCase()  +"@"+  CLUSTER_NAME;
          });
  
          when(mockContext.getQualifiedNameForColumn(any(LineageVertex.class))).thenAnswer(invocation -> {
              LineageVertex vertex = invocation.getArgument(0);
              return vertex.getVertexId().toLowerCase()  +"@"+  CLUSTER_NAME;
          });
  
          when(mockContext.getDatabaseNameFromTable(anyString())).thenAnswer(invocation -> {
              String tableName = invocation.getArgument(0);
              if (tableName.contains(".")) {
                  return tableName.split("\\.")[0];
              }
              return DB_NAME;
          });
  
          when(mockContext.getColumnNameOnly(anyString())).thenAnswer(invocation -> {
              String columnName = invocation.getArgument(0);
              if (columnName.contains(".")) {
                  String[] parts = columnName.split("\\.");
                  return parts[parts.length - 1];
              }
              return columnName;
          });
  
          when(mockContext.getTableNameFromColumn(anyString())).thenAnswer(invocation -> {
              String columnName = invocation.getArgument(0);
              if (columnName.contains(".") && columnName.split("\\.").length >= 2) {
                  String[] parts = columnName.split("\\.");
                  return parts[0]  +"."+  parts[1];
              }
              return DB_NAME  +"."+  TABLE_NAME;
          });
      }
  
      @Test
      public void testGetUserName() {
          assertEquals(baseImpalaEvent.getUserName(), USER_NAME);
      }
  
      @Test
      public void testGetQualifiedNameWithImpalaNode() throws Exception {
          // Create a test ImpalaNode with database vertex
          LineageVertex dbVertex = createDatabaseVertex(DB_NAME);
          ImpalaNode dbNode = new ImpalaNode(dbVertex);
  
          String qualifiedName = baseImpalaEvent.getQualifiedName(dbNode);
  
          assertNotNull(qualifiedName);
          assertEquals(qualifiedName, DB_NAME.toLowerCase()  +"@"+  CLUSTER_NAME);
      }
  
      @Test
      public void testGetQualifiedNameWithLineageVertex() throws Exception {
          // Test database vertex
          LineageVertex dbVertex = createDatabaseVertex(DB_NAME);
          String dbQualifiedName = baseImpalaEvent.getQualifiedName(dbVertex);
          assertEquals(dbQualifiedName, DB_NAME.toLowerCase()  +"@"+  CLUSTER_NAME);
  
          // Test table vertex
          LineageVertex tableVertex = createTableVertex(DB_NAME+"."+TABLE_NAME);
          String tableQualifiedName = baseImpalaEvent.getQualifiedName(tableVertex);
          assertEquals(tableQualifiedName, DB_NAME+"."+TABLE_NAME.toLowerCase()  +"@"+  CLUSTER_NAME);
  
          // Test column vertex
          LineageVertex columnVertex = createColumnVertex(COLUMN_NAME, TABLE_NAME);
          String columnQualifiedName = baseImpalaEvent.getQualifiedName(columnVertex);
          System.out.println("columnQualifiedName : "+columnQualifiedName);
          assertEquals(columnQualifiedName, DB_NAME+"."+TABLE_NAME+"."+COLUMN_NAME.toLowerCase()  +"@"+  CLUSTER_NAME);
      }
  
      @Test
      public void testGetQualifiedNameWithNullVertex() {
          try {
              baseImpalaEvent.getQualifiedName((LineageVertex) null);
              fail("Expected IllegalArgumentException for null vertex");
          } catch (IllegalArgumentException e) {
              assertEquals(e.getMessage(), "node is null");
          }
      }
  
      @Test
      public void testGetQualifiedNameWithNullVertexType() {
          LineageVertex vertex = new LineageVertex();
          vertex.setVertexId("test");
          vertex.setVertexType(null);
  
          String qualifiedName = baseImpalaEvent.getQualifiedName(vertex);
          assertNull(qualifiedName);
      }
  
      @Test
      public void testGetQualifiedNameWithNullVertexId() {
          LineageVertex vertex = new LineageVertex();
          vertex.setVertexType(ImpalaVertexType.DATABASE);
          vertex.setVertexId(null);
  
          String qualifiedName = baseImpalaEvent.getQualifiedName(vertex);
          assertNull(qualifiedName);
      }
  
      @Test
      public void testGetTableNameFromVertex() {
          // Test with column vertex that has metadata
          LineageVertex columnVertex = createColumnVertex(COLUMN_NAME, TABLE_NAME);
          LineageVertexMetadata metadata = new LineageVertexMetadata();
          metadata.setTableName(TABLE_NAME);
          columnVertex.setMetadata(metadata);
  
          String tableName = baseImpalaEvent.getTableNameFromVertex(columnVertex);
          assertEquals(tableName, TABLE_NAME);
  
          // Test with non-column vertex
          LineageVertex dbVertex = createDatabaseVertex(DB_NAME);
          String tableNameFromDb = baseImpalaEvent.getTableNameFromVertex(dbVertex);
          assertEquals(tableNameFromDb, DB_NAME  +"."+  TABLE_NAME);
      }
  
      @Test
      public void testGetTableNameFromColumn() {
          String columnName = DB_NAME  +"."+  TABLE_NAME  +"."+  COLUMN_NAME;
          String tableName = baseImpalaEvent.getTableNameFromColumn(columnName);
          assertEquals(tableName, DB_NAME  +"."+  TABLE_NAME);
      }
  
      @Test
      public void testToDbEntityWithString() throws Exception {
  
          when(mockContext.getEntity(anyString())).thenReturn(null);
  
          AtlasEntity dbEntity = baseImpalaEvent.toDbEntity(DB_NAME);
  
          assertNotNull(dbEntity);
          assertEquals(dbEntity.getTypeName(), HIVE_TYPE_DB);
          assertEquals(dbEntity.getAttribute(ATTRIBUTE_NAME), DB_NAME.toLowerCase());
          assertEquals(dbEntity.getAttribute(ATTRIBUTE_QUALIFIED_NAME), DB_NAME.toLowerCase()  +"@"+  CLUSTER_NAME);
          assertEquals(dbEntity.getAttribute(ATTRIBUTE_CLUSTER_NAME), CLUSTER_NAME);
          assertNull(dbEntity.getGuid());
      }
  
  
      @Test
      public void testGetColumnEntities() throws Exception {
  
          LineageVertex tableVertex = createTableVertex(DB_NAME+"."+TABLE_NAME);
          tableVertex.setId(1L);
          tableVertex.setCreateTime(TIMESTAMP);
          ImpalaNode tableNode = new ImpalaNode(tableVertex);
          tableNode.addChild(tableVertex);
          AtlasEntity entity = new AtlasEntity(HIVE_TYPE_COLUMN);
          entity.setGuid("test-guid");
          entity.setAttribute(ATTRIBUTE_QUALIFIED_NAME, "test.qualified.name");
  
          AtlasObjectId objectId = BaseImpalaEvent.getObjectId(entity);
  
          when(mockContext.getEntity(anyString())).thenReturn(null);
          when(mockContext.getUserName()).thenReturn(USER_NAME);
  
          List<AtlasEntity> entityList = baseImpalaEvent.getColumnEntities(objectId ,tableNode);
  
          assertNotNull(entityList.get(0));
          assertEquals(entityList.get(0).getTypeName(), HIVE_TYPE_COLUMN);
          assertEquals(entityList.get(0).getAttribute(ATTRIBUTE_TABLE),objectId);
      }
  
      @Test
      public void testToTableEntity() throws Exception {
  
          LineageVertex tableVertex = createTableVertex(DB_NAME+"."+TABLE_NAME);
          tableVertex.setId(1L);
          tableVertex.setCreateTime(TIMESTAMP);
          ImpalaNode tableNode = new ImpalaNode(tableVertex);
          tableNode.addChild(tableVertex);
          AtlasEntity entity = new AtlasEntity(HIVE_TYPE_TABLE);
          entity.setGuid("test-guid");
  
          AtlasObjectId objectId = BaseImpalaEvent.getObjectId(entity);
  
          when(mockContext.getEntity(anyString())).thenReturn(null);
          when(mockContext.getUserName()).thenReturn(USER_NAME);
  
          AtlasEntity atlasEntity = baseImpalaEvent.toTableEntity(objectId ,tableNode,null);
  
  
          assertEquals(atlasEntity.getTypeName(),HIVE_TYPE_TABLE);
          assertEquals(atlasEntity.getAttribute(ATTRIBUTE_QUALIFIED_NAME),DB_NAME+"."+TABLE_NAME+"@"+CLUSTER_NAME);
      }
  
  
  
      @Test
      public void testToDbEntityWithCachedEntity() throws Exception {
          // Setup entity cache to return existing entity
          AtlasEntity cachedEntity = new AtlasEntity(HIVE_TYPE_DB);
          cachedEntity.setAttribute(ATTRIBUTE_NAME, DB_NAME);
          when(mockContext.getEntity(anyString())).thenReturn(cachedEntity);
  
          AtlasEntity dbEntity = baseImpalaEvent.toDbEntity(DB_NAME);
  
          assertNotNull(dbEntity);
          assertEquals(dbEntity, cachedEntity);
      }
  
  
      @Test
      public void testToTableEntityWithNullNode() {
          try {
              baseImpalaEvent.toTableEntity((ImpalaNode) null, (AtlasEntitiesWithExtInfo) null);
              fail("Expected IllegalArgumentException for null table node");
          } catch (Exception e) {
              assertTrue(e instanceof IllegalArgumentException);
              assertTrue(e.getMessage().contains("table is null"));
          }
      }
  
      @Test
      public void testGetObjectId() {
          AtlasEntity entity = new AtlasEntity(HIVE_TYPE_TABLE);
          entity.setGuid("test-guid");
          entity.setAttribute(ATTRIBUTE_QUALIFIED_NAME, "test.qualified.name");
  
          AtlasObjectId objectId = BaseImpalaEvent.getObjectId(entity);
  
          assertNotNull(objectId);
          assertEquals(objectId.getGuid(), "test-guid");
          assertEquals(objectId.getTypeName(), HIVE_TYPE_TABLE);
          assertEquals(objectId.getUniqueAttributes().get(ATTRIBUTE_QUALIFIED_NAME), "test.qualified.name");
      }
  
      @Test
      public void testGetObjectIds() {
          List<AtlasEntity> entities = new ArrayList<>();
  
          AtlasEntity entity1 = new AtlasEntity(HIVE_TYPE_TABLE);
          entity1.setGuid("test-guid-1");
          entity1.setAttribute(ATTRIBUTE_QUALIFIED_NAME, "test.table1");
          entities.add(entity1);
  
          AtlasEntity entity2 = new AtlasEntity(HIVE_TYPE_COLUMN);
          entity2.setGuid("test-guid-2");
          entity2.setAttribute(ATTRIBUTE_QUALIFIED_NAME, "test.column1");
          entities.add(entity2);
  
          List<AtlasObjectId> objectIds = BaseImpalaEvent.getObjectIds(entities);
  
          assertNotNull(objectIds);
          assertEquals(objectIds.size(), 2);
          assertEquals(objectIds.get(0).getGuid(), "test-guid-1");
          assertEquals(objectIds.get(1).getGuid(), "test-guid-2");
      }
  
      @Test
      public void testGetObjectIdsWithEmptyList() {
          List<AtlasObjectId> objectIds = BaseImpalaEvent.getObjectIds(Collections.emptyList());
          assertNotNull(objectIds);
          assertTrue(objectIds.isEmpty());
      }
  
      @Test
      public void testGetTableCreateTimeWithNode() {
          LineageVertex tableVertex = createTableVertex(DB_NAME+"."+TABLE_NAME);
          tableVertex.setCreateTime(TIMESTAMP);
          ImpalaNode tableNode = new ImpalaNode(tableVertex);
  
          long createTime = BaseImpalaEvent.getTableCreateTime(tableNode);
          assertEquals(createTime, TIMESTAMP * MILLIS_CONVERT_FACTOR);
      }
  
      @Test
      public void testGetTableCreateTimeWithVertex() {
          LineageVertex tableVertex = createTableVertex(DB_NAME+"."+TABLE_NAME);
          tableVertex.setCreateTime(TIMESTAMP);
  
          long createTime = BaseImpalaEvent.getTableCreateTime(tableVertex);
          assertEquals(createTime, TIMESTAMP * MILLIS_CONVERT_FACTOR);
      }
  
      @Test
      public void testGetTableCreateTimeWithNullTime() {
          LineageVertex tableVertex = createTableVertex(DB_NAME+"."+TABLE_NAME);
          tableVertex.setCreateTime(null);
  
          long createTime = BaseImpalaEvent.getTableCreateTime(tableVertex);
          assertTrue(createTime > 0);
      }
  
      @Test
      public void testGetImpalaProcessEntity() throws Exception {
          List<AtlasEntity> inputs = createMockEntities("input_table");
          List<AtlasEntity> outputs = createMockEntities("output_table");
  
          AtlasEntity processEntity = baseImpalaEvent.getImpalaProcessEntity(inputs, outputs);
  
          assertNotNull(processEntity);
          assertEquals(processEntity.getTypeName(), ImpalaDataType.IMPALA_PROCESS.getName());
          assertNotNull(processEntity.getAttribute(ATTRIBUTE_QUALIFIED_NAME));
          assertNotNull(processEntity.getAttribute(ATTRIBUTE_INPUTS));
          assertNotNull(processEntity.getAttribute(ATTRIBUTE_OUTPUTS));
          assertEquals(processEntity.getAttribute(ATTRIBUTE_OPERATION_TYPE), ImpalaOperationType.QUERY);
          assertEquals(processEntity.getAttribute(ATTRIBUTE_START_TIME), TIMESTAMP * MILLIS_CONVERT_FACTOR);
          assertEquals(processEntity.getAttribute(ATTRIBUTE_END_TIME), END_TIME * MILLIS_CONVERT_FACTOR);
          assertEquals(processEntity.getAttribute(ATTRIBUTE_USER_NAME), EMPTY_ATTRIBUTE_VALUE);
          assertEquals(processEntity.getAttribute(ATTRIBUTE_QUERY_TEXT), EMPTY_ATTRIBUTE_VALUE);
      }
  
      @Test
      public void testGetImpalaProcessExecutionEntity() throws Exception {
          List<AtlasEntity> inputs = createMockEntities("input_table");
          List<AtlasEntity> outputs = createMockEntities("output_table");
          AtlasEntity processEntity = baseImpalaEvent.getImpalaProcessEntity(inputs, outputs);
  
          AtlasEntity executionEntity = baseImpalaEvent.getImpalaProcessExecutionEntity(processEntity);
  
          assertNotNull(executionEntity);
          assertEquals(executionEntity.getTypeName(), ImpalaDataType.IMPALA_PROCESS_EXECUTION.getName());
          assertNotNull(executionEntity.getAttribute(ATTRIBUTE_QUALIFIED_NAME));
          assertEquals(executionEntity.getAttribute(ATTRIBUTE_START_TIME), TIMESTAMP * MILLIS_CONVERT_FACTOR);
          assertEquals(executionEntity.getAttribute(ATTRIBUTE_END_TIME), END_TIME * MILLIS_CONVERT_FACTOR);
          assertEquals(executionEntity.getAttribute(ATTRIBUTE_USER_NAME), USER_NAME);
          assertEquals(executionEntity.getAttribute(ATTRIBUTE_QUERY_TEXT), QUERY_TEXT.toLowerCase().trim());
          assertEquals(executionEntity.getAttribute(ATTRIBUTE_QUERY_ID), QUERY_ID);
          assertEquals(executionEntity.getAttribute(ATTRIBUTE_HOSTNAME), HOST_NAME);
      }
  
      @Test
      public void testCreateTableNode() {
          Long createTime = TIMESTAMP;
          ImpalaNode tableNode = baseImpalaEvent.createTableNode(TABLE_NAME, createTime);
  
          assertNotNull(tableNode);
          assertEquals(tableNode.getOwnVertex().getVertexType(), ImpalaVertexType.TABLE);
          assertEquals(tableNode.getOwnVertex().getVertexId(), TABLE_NAME);
          assertEquals(tableNode.getOwnVertex().getCreateTime(), createTime);
      }
  
      @Test
      public void testCreateHiveDDLEntityWithDb() {
          AtlasEntity dbEntity = new AtlasEntity(HIVE_TYPE_DB);
          dbEntity.setAttribute(ATTRIBUTE_QUALIFIED_NAME, "test.db.qualified.name");
  
          AtlasEntity ddlEntity = baseImpalaEvent.createHiveDDLEntity(dbEntity);
  
          assertNotNull(ddlEntity);
          assertEquals(ddlEntity.getTypeName(), ImpalaDataType.HIVE_DB_DDL.getName());
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_SERVICE_TYPE), "impala");
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_EXEC_TIME), TIMESTAMP * MILLIS_CONVERT_FACTOR);
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_QUERY_TEXT), QUERY_TEXT);
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_USER_NAME), USER_NAME);
          assertNotNull(ddlEntity.getAttribute(ATTRIBUTE_DB));
      }
  
      @Test
      public void testCreateHiveDDLEntityWithTable() {
          AtlasEntity tableEntity = new AtlasEntity(HIVE_TYPE_TABLE);
          tableEntity.setAttribute(ATTRIBUTE_QUALIFIED_NAME, "test.table.qualified.name");
  
          AtlasEntity ddlEntity = baseImpalaEvent.createHiveDDLEntity(tableEntity);
  
          assertNotNull(ddlEntity);
          assertEquals(ddlEntity.getTypeName(), ImpalaDataType.HIVE_TABLE_DDL.getName());
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_SERVICE_TYPE), "impala");
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_EXEC_TIME), TIMESTAMP * MILLIS_CONVERT_FACTOR);
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_QUERY_TEXT), QUERY_TEXT);
          assertEquals(ddlEntity.getAttribute(ATTRIBUTE_USER_NAME), USER_NAME);
          assertNotNull(ddlEntity.getAttribute(ATTRIBUTE_TABLE));
      }
  
      @Test
      public void testIsDdlOperation() {
  
          when(mockContext.getImpalaOperationType()).thenReturn(ImpalaOperationType.CREATEVIEW);
          assertTrue(baseImpalaEvent.isDdlOperation());
  
          when(mockContext.getImpalaOperationType()).thenReturn(ImpalaOperationType.ALTERVIEW_AS);
          assertTrue(baseImpalaEvent.isDdlOperation());
  
          when(mockContext.getImpalaOperationType()).thenReturn(ImpalaOperationType.CREATETABLE_AS_SELECT);
          assertTrue(baseImpalaEvent.isDdlOperation());
  
          when(mockContext.getImpalaOperationType()).thenReturn(ImpalaOperationType.QUERY);
          assertFalse(baseImpalaEvent.isDdlOperation());
      }
  
      @Test
      public void testGetCreateTimeInVertex() {
  
          LineageVertex vertex = createDatabaseVertex(DB_NAME);
          vertex.setCreateTime(TIMESTAMP);
  
          Long createTime = baseImpalaEvent.getCreateTimeInVertex(vertex);
          assertEquals(createTime, Long.valueOf(TIMESTAMP));
  
          Long createTimeNull = baseImpalaEvent.getCreateTimeInVertex(null);
          assertTrue(createTimeNull > 0); // Should return current time in seconds
  
          LineageVertex columnVertex = createColumnVertex(COLUMN_NAME, TABLE_NAME);
          columnVertex.setCreateTime(null);
          LineageVertexMetadata metadata = new LineageVertexMetadata();
          metadata.setTableCreateTime(TIMESTAMP);
          columnVertex.setMetadata(metadata);
  
          Long metadataCreateTime = baseImpalaEvent.getCreateTimeInVertex(columnVertex);
          assertEquals(metadataCreateTime, Long.valueOf(TIMESTAMP));
      }
  
      @Test(dataProvider = "qualifiedNameDataProvider")
      public void testGetQualifiedNameForProcess(ImpalaOperationType operationType,
                                                 List<AtlasEntity> inputs,
                                                 List<AtlasEntity> outputs,
                                                 boolean shouldThrowException) throws Exception {
          when(mockContext.getImpalaOperationType()).thenReturn(operationType);
  
          if (shouldThrowException) {
              try {
                  baseImpalaEvent.getQualifiedName(inputs, outputs);
                  fail("Expected IllegalArgumentException");
              } catch (IllegalArgumentException e) {
                  assertTrue(e.getMessage().contains("unexpected operation type"));
              }
          } else {
              String qualifiedName = baseImpalaEvent.getQualifiedName(inputs, outputs);
              assertNotNull(qualifiedName);
              if (operationType == ImpalaOperationType.CREATEVIEW ||
                      operationType == ImpalaOperationType.CREATETABLE_AS_SELECT ||
                      operationType == ImpalaOperationType.ALTERVIEW_AS) {
                  assertTrue(qualifiedName.contains("@"));
              } else {
                  assertTrue(qualifiedName.contains("->"));
              }
          }
      }
  
      @DataProvider(name = "qualifiedNameDataProvider")
      public Object[][] qualifiedNameDataProvider() {
          List<AtlasEntity> inputs = createMockEntities("input_table");
          List<AtlasEntity> outputs = createMockEntities("output_table");
  
          return new Object[][] {
                  {ImpalaOperationType.CREATEVIEW, inputs, outputs, false},
                  {ImpalaOperationType.CREATETABLE_AS_SELECT, inputs, outputs, false},
                  {ImpalaOperationType.ALTERVIEW_AS, inputs, outputs, false},
                  {ImpalaOperationType.QUERY, inputs, outputs, false},
                  {ImpalaOperationType.QUERY_WITH_CLAUSE, inputs, outputs, false},
                  {ImpalaOperationType.INSERT, inputs, outputs, true}  // Should throw exception
          };
      }
  
      // Helper methods for creating test data
  
      private LineageVertex createDatabaseVertex(String dbName) {
          LineageVertex vertex = new LineageVertex();
          vertex.setVertexType(ImpalaVertexType.DATABASE);
          vertex.setVertexId(dbName);
          vertex.setCreateTime(TIMESTAMP);
          return vertex;
      }
  
      private LineageVertex createTableVertex(String tableName) {
          LineageVertex vertex = new LineageVertex();
          vertex.setVertexType(ImpalaVertexType.TABLE);
          vertex.setVertexId(tableName);
          vertex.setCreateTime(TIMESTAMP);
          return vertex;
      }
  
      private LineageVertex createColumnVertex(String columnName, String tableName) {
          LineageVertex vertex = new LineageVertex();
          vertex.setVertexType(ImpalaVertexType.COLUMN);
          vertex.setVertexId(DB_NAME+"."+tableName+"."+columnName);
          vertex.setCreateTime(TIMESTAMP);
  
          LineageVertexMetadata metadata = new LineageVertexMetadata();
          metadata.setTableName(tableName);
          metadata.setTableCreateTime(TIMESTAMP);
          vertex.setMetadata(metadata);
  
          return vertex;
      }
  
      private List<AtlasEntity> createMockEntities(String tableBaseName) {
          List<AtlasEntity> entities = new ArrayList<>();
  
          AtlasEntity entity = new AtlasEntity(HIVE_TYPE_TABLE);
          entity.setGuid("test-guid-" + tableBaseName);
          entity.setAttribute(ATTRIBUTE_QUALIFIED_NAME, tableBaseName  +"@"+  CLUSTER_NAME);
          entity.setAttribute(ATTRIBUTE_NAME, tableBaseName);
          entity.setAttribute(ATTRIBUTE_CREATE_TIME, TIMESTAMP * MILLIS_CONVERT_FACTOR);
  
          entities.add(entity);
          return entities;
      }
  }
  