/*
 * 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.gravitino.iceberg.service.dispatcher;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import java.io.IOException;
import java.time.Instant;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.apache.gravitino.Config;
import org.apache.gravitino.Configs;
import org.apache.gravitino.Entity;
import org.apache.gravitino.EntityStore;
import org.apache.gravitino.GravitinoEnv;
import org.apache.gravitino.NameIdentifier;
import org.apache.gravitino.authorization.Owner;
import org.apache.gravitino.authorization.OwnerDispatcher;
import org.apache.gravitino.catalog.TableDispatcher;
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.iceberg.common.utils.IcebergIdentifierUtils;
import org.apache.gravitino.iceberg.service.authorization.IcebergRESTServerContext;
import org.apache.gravitino.iceberg.service.provider.IcebergConfigProvider;
import org.apache.gravitino.listener.api.event.IcebergRequestContext;
import org.apache.gravitino.meta.AuditInfo;
import org.apache.gravitino.meta.TableEntity;
import org.apache.iceberg.Schema;
import org.apache.iceberg.UpdateRequirement;
import org.apache.iceberg.catalog.Namespace;
import org.apache.iceberg.catalog.TableIdentifier;
import org.apache.iceberg.rest.requests.CreateTableRequest;
import org.apache.iceberg.rest.requests.RenameTableRequest;
import org.apache.iceberg.rest.requests.UpdateTableRequest;
import org.apache.iceberg.rest.responses.LoadTableResponse;
import org.apache.iceberg.types.Types.NestedField;
import org.apache.iceberg.types.Types.StringType;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

public class TestIcebergTableHookDispatcher {

  private static final String TEST_METALAKE = "test_metalake";
  private static final String TEST_CATALOG = "test_catalog";
  private static final String TEST_USER = "test_user";
  private static final Schema TABLE_SCHEMA =
      new Schema(NestedField.required(1, "test_field", StringType.get()));

  private IcebergTableHookDispatcher hookDispatcher;
  private IcebergTableOperationDispatcher mockDispatcher;
  private IcebergNamespaceOperationDispatcher mockNamespaceDispatcher;
  private EntityStore mockEntityStore;
  private TableDispatcher mockTableDispatcher;
  private TableDispatcher mockInternalTableDispatcher;
  private OwnerDispatcher mockOwnerDispatcher;
  private OwnerDispatcher mockInternalOwnerDispatcher;
  private IcebergRequestContext mockContext;

  private Config previousConfig;
  private EntityStore previousEntityStore;
  private TableDispatcher previousTableDispatcher;
  private TableDispatcher previousInternalTableDispatcher;
  private OwnerDispatcher previousOwnerDispatcher;
  private OwnerDispatcher previousInternalOwnerDispatcher;

  @BeforeEach
  public void setUp() throws IllegalAccessException {
    // Mock the underlying dispatcher
    mockDispatcher = mock(IcebergTableOperationDispatcher.class);
    mockNamespaceDispatcher = mock(IcebergNamespaceOperationDispatcher.class);

    // Mock GravitinoEnv components
    mockEntityStore = mock(EntityStore.class);
    mockTableDispatcher = mock(TableDispatcher.class);
    mockInternalTableDispatcher = mock(TableDispatcher.class);
    mockOwnerDispatcher = mock(OwnerDispatcher.class);
    mockInternalOwnerDispatcher = mock(OwnerDispatcher.class);

    Config mockConfig = mock(Config.class);
    when(mockConfig.get(Configs.SCHEMA_SEPARATOR)).thenReturn(":");
    previousConfig = (Config) FieldUtils.readField(GravitinoEnv.getInstance(), "config", true);
    previousEntityStore =
        (EntityStore) FieldUtils.readField(GravitinoEnv.getInstance(), "entityStore", true);
    previousTableDispatcher =
        (TableDispatcher) FieldUtils.readField(GravitinoEnv.getInstance(), "tableDispatcher", true);
    previousInternalTableDispatcher =
        (TableDispatcher)
            FieldUtils.readField(GravitinoEnv.getInstance(), "internalTableDispatcher", true);
    previousOwnerDispatcher =
        (OwnerDispatcher) FieldUtils.readField(GravitinoEnv.getInstance(), "ownerDispatcher", true);
    previousInternalOwnerDispatcher =
        (OwnerDispatcher)
            FieldUtils.readField(GravitinoEnv.getInstance(), "internalOwnerDispatcher", true);
    FieldUtils.writeField(GravitinoEnv.getInstance(), "config", mockConfig, true);
    FieldUtils.writeField(GravitinoEnv.getInstance(), "entityStore", mockEntityStore, true);
    FieldUtils.writeField(GravitinoEnv.getInstance(), "tableDispatcher", mockTableDispatcher, true);
    FieldUtils.writeField(
        GravitinoEnv.getInstance(), "internalTableDispatcher", mockInternalTableDispatcher, true);
    FieldUtils.writeField(GravitinoEnv.getInstance(), "ownerDispatcher", mockOwnerDispatcher, true);
    FieldUtils.writeField(
        GravitinoEnv.getInstance(), "internalOwnerDispatcher", mockInternalOwnerDispatcher, true);

    // Create mock IcebergRESTServerContext
    IcebergConfigProvider mockConfigProvider = mock(IcebergConfigProvider.class);
    when(mockConfigProvider.getMetalakeName()).thenReturn(TEST_METALAKE);
    when(mockConfigProvider.getDefaultCatalogName()).thenReturn(TEST_CATALOG);
    IcebergRESTServerContext.create(mockConfigProvider, false, false, true, null);

    // Create hook dispatcher
    hookDispatcher = new IcebergTableHookDispatcher(mockDispatcher, mockNamespaceDispatcher);

    // Mock request context
    mockContext = mock(IcebergRequestContext.class);
    when(mockContext.catalogName()).thenReturn(TEST_CATALOG);
    when(mockContext.userName()).thenReturn(TEST_USER);
  }

  @AfterEach
  public void tearDown() throws IllegalAccessException {
    // Restore GravitinoEnv to its prior state so this test doesn't leak singleton state
    // into other tests that share the same JVM.
    FieldUtils.writeField(GravitinoEnv.getInstance(), "config", previousConfig, true);
    FieldUtils.writeField(GravitinoEnv.getInstance(), "entityStore", previousEntityStore, true);
    FieldUtils.writeField(
        GravitinoEnv.getInstance(), "tableDispatcher", previousTableDispatcher, true);
    FieldUtils.writeField(
        GravitinoEnv.getInstance(),
        "internalTableDispatcher",
        previousInternalTableDispatcher,
        true);
    FieldUtils.writeField(
        GravitinoEnv.getInstance(), "ownerDispatcher", previousOwnerDispatcher, true);
    FieldUtils.writeField(
        GravitinoEnv.getInstance(),
        "internalOwnerDispatcher",
        previousInternalOwnerDispatcher,
        true);

    // Reset IcebergRESTServerContext singleton
    Class<?> holderClass =
        Arrays.stream(IcebergRESTServerContext.class.getDeclaredClasses())
            .filter(c -> c.getSimpleName().equals("InstanceHolder"))
            .findFirst()
            .orElseThrow(() -> new RuntimeException("InstanceHolder class not found"));
    FieldUtils.writeStaticField(holderClass, "INSTANCE", null, true);
  }

  @Test
  public void testCreateTableCallsImportAndOwnership() {
    Namespace namespace = Namespace.of("test_schema");
    CreateTableRequest request =
        CreateTableRequest.builder().withName("test_table").withSchema(TABLE_SCHEMA).build();

    LoadTableResponse mockResponse = mock(LoadTableResponse.class);
    when(mockDispatcher.createTable(mockContext, namespace, request)).thenReturn(mockResponse);

    LoadTableResponse result = hookDispatcher.createTable(mockContext, namespace, request);

    Assertions.assertEquals(mockResponse, result);
    verify(mockDispatcher).createTable(mockContext, namespace, request);

    // Verify table import was called
    NameIdentifier expectedIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            TEST_METALAKE, TEST_CATALOG, TableIdentifier.of(namespace, "test_table"), ":");
    verify(mockInternalTableDispatcher).loadTable(expectedIdentifier);
    verify(mockTableDispatcher, never()).loadTable(any());

    // Verify ownership was set
    ArgumentCaptor<String> userCaptor = ArgumentCaptor.forClass(String.class);
    verify(mockInternalOwnerDispatcher)
        .setOwner(eq(TEST_METALAKE), any(), userCaptor.capture(), eq(Owner.Type.USER));
    verify(mockOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
    Assertions.assertEquals(TEST_USER, userCaptor.getValue());
  }

  @Test
  public void testDropTableDeletesEntity() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");

    hookDispatcher.dropTable(mockContext, tableId, false);

    verify(mockDispatcher).dropTable(mockContext, tableId, false);

    NameIdentifier expectedIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            TEST_METALAKE, TEST_CATALOG, tableId, ":");
    verify(mockEntityStore).delete(expectedIdentifier, Entity.EntityType.TABLE);
  }

  @Test
  public void testDropTableReimportsEntityWhenTableExistsAfterDrop() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");
    when(mockDispatcher.tableExists(mockContext, tableId)).thenReturn(true);

    hookDispatcher.dropTable(mockContext, tableId, false);

    verify(mockDispatcher).dropTable(mockContext, tableId, false);
    NameIdentifier expectedIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            TEST_METALAKE, TEST_CATALOG, tableId, ":");
    verify(mockEntityStore, never()).delete(expectedIdentifier, Entity.EntityType.TABLE);
    verify(mockInternalTableDispatcher).loadTable(expectedIdentifier);
    verify(mockTableDispatcher, never()).loadTable(any());
  }

  @Test
  public void testDropTableReimportsEntityWhenTableIsRecreatedDuringDelete() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");
    when(mockDispatcher.tableExists(mockContext, tableId)).thenReturn(false, true);

    hookDispatcher.dropTable(mockContext, tableId, false);

    verify(mockDispatcher).dropTable(mockContext, tableId, false);
    NameIdentifier expectedIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            TEST_METALAKE, TEST_CATALOG, tableId, ":");
    verify(mockEntityStore).delete(expectedIdentifier, Entity.EntityType.TABLE);
    verify(mockInternalTableDispatcher).loadTable(expectedIdentifier);
    verify(mockTableDispatcher, never()).loadTable(any());
  }

  @Test
  public void testDropTableCleansUpOrphanedSchemaEntities() throws IOException {
    TableIdentifier tableId = TableIdentifier.of(Namespace.of("parent", "child"), "test_table");
    NameIdentifier schemaIdent = NameIdentifier.of(TEST_METALAKE, TEST_CATALOG, "parent");
    when(mockNamespaceDispatcher.namespaceExists(mockContext, Namespace.of("parent", "child")))
        .thenReturn(false);
    when(mockNamespaceDispatcher.namespaceExists(mockContext, Namespace.of("parent")))
        .thenReturn(false);

    hookDispatcher.dropTable(mockContext, tableId, false);

    verify(mockDispatcher).dropTable(mockContext, tableId, false);
    verify(mockEntityStore).delete(schemaIdent, Entity.EntityType.SCHEMA, true);
  }

  @Test
  public void testDropTableIgnoresNoSuchEntityException() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");

    NameIdentifier expectedIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            TEST_METALAKE, TEST_CATALOG, tableId, ":");
    doThrow(new NoSuchEntityException("Table not found"))
        .when(mockEntityStore)
        .delete(expectedIdentifier, Entity.EntityType.TABLE);

    // Should not throw exception
    Assertions.assertDoesNotThrow(() -> hookDispatcher.dropTable(mockContext, tableId, false));

    verify(mockDispatcher).dropTable(mockContext, tableId, false);
  }

  @Test
  public void testDropTableIgnoresReconciliationIOException() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");

    doThrow(new IOException("IO error")).when(mockEntityStore).delete(any(), any());

    Assertions.assertDoesNotThrow(() -> hookDispatcher.dropTable(mockContext, tableId, false));
    verify(mockDispatcher).dropTable(mockContext, tableId, false);
  }

  @Test
  public void testRenameTableUpdatesEntity() throws IOException {
    TableIdentifier source = TableIdentifier.of("schema1", "old_table");
    TableIdentifier dest = TableIdentifier.of("schema2", "new_table");
    RenameTableRequest request =
        RenameTableRequest.builder().withSource(source).withDestination(dest).build();

    TableEntity mockTableEntity = mock(TableEntity.class);
    when(mockTableEntity.id()).thenReturn(1L);
    when(mockTableEntity.columns()).thenReturn(Collections.emptyList());
    AuditInfo auditInfo =
        AuditInfo.builder().withCreator("original_creator").withCreateTime(Instant.now()).build();
    when(mockTableEntity.auditInfo()).thenReturn(auditInfo);

    when(mockEntityStore.update(any(), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any()))
        .thenReturn(mockTableEntity);

    hookDispatcher.renameTable(mockContext, request);

    verify(mockDispatcher).renameTable(mockContext, request);

    NameIdentifier sourceIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(TEST_METALAKE, TEST_CATALOG, source, ":");
    verify(mockEntityStore)
        .update(eq(sourceIdentifier), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any());
  }

  @Test
  public void testRenameTableReconcilesSourceAndDestinationEntities() throws IOException {
    TableIdentifier source = TableIdentifier.of("schema1", "old_table");
    TableIdentifier dest = TableIdentifier.of("schema2", "new_table");
    RenameTableRequest request =
        RenameTableRequest.builder().withSource(source).withDestination(dest).build();

    TableEntity mockTableEntity = mock(TableEntity.class);
    when(mockTableEntity.id()).thenReturn(1L);
    when(mockTableEntity.columns()).thenReturn(Collections.emptyList());
    AuditInfo auditInfo =
        AuditInfo.builder().withCreator("original_creator").withCreateTime(Instant.now()).build();
    when(mockTableEntity.auditInfo()).thenReturn(auditInfo);
    when(mockEntityStore.update(any(), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any()))
        .thenReturn(mockTableEntity);
    when(mockDispatcher.tableExists(mockContext, source)).thenReturn(false, false);
    when(mockDispatcher.tableExists(mockContext, dest)).thenReturn(true);

    hookDispatcher.renameTable(mockContext, request);

    NameIdentifier sourceIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(TEST_METALAKE, TEST_CATALOG, source, ":");
    NameIdentifier destIdentifier =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(TEST_METALAKE, TEST_CATALOG, dest, ":");
    verify(mockEntityStore).delete(sourceIdentifier, Entity.EntityType.TABLE);
    verify(mockInternalTableDispatcher).loadTable(destIdentifier);
    verify(mockTableDispatcher, never()).loadTable(any());
  }

  @Test
  public void testRenameTableIgnoresReconciliationImportFailure() throws IOException {
    TableIdentifier source = TableIdentifier.of("schema1", "old_table");
    TableIdentifier dest = TableIdentifier.of("schema2", "new_table");
    RenameTableRequest request =
        RenameTableRequest.builder().withSource(source).withDestination(dest).build();

    TableEntity mockTableEntity = mock(TableEntity.class);
    when(mockTableEntity.id()).thenReturn(1L);
    when(mockTableEntity.columns()).thenReturn(Collections.emptyList());
    AuditInfo auditInfo =
        AuditInfo.builder().withCreator("original_creator").withCreateTime(Instant.now()).build();
    when(mockTableEntity.auditInfo()).thenReturn(auditInfo);
    when(mockEntityStore.update(any(), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any()))
        .thenReturn(mockTableEntity);
    when(mockDispatcher.tableExists(mockContext, dest)).thenReturn(true);
    doThrow(new RuntimeException("import failed"))
        .when(mockInternalTableDispatcher)
        .loadTable(any());

    Assertions.assertDoesNotThrow(() -> hookDispatcher.renameTable(mockContext, request));

    verify(mockDispatcher).renameTable(mockContext, request);
    verify(mockInternalTableDispatcher).loadTable(any());
    verify(mockTableDispatcher, never()).loadTable(any());
  }

  @Test
  public void testRenameTableIgnoresNoSuchEntityException() throws IOException {
    TableIdentifier source = TableIdentifier.of("schema1", "old_table");
    TableIdentifier dest = TableIdentifier.of("schema2", "new_table");
    RenameTableRequest request =
        RenameTableRequest.builder().withSource(source).withDestination(dest).build();

    doThrow(new NoSuchEntityException("Entity not found"))
        .when(mockEntityStore)
        .update(any(), any(), any(), any());

    // Should not throw exception
    Assertions.assertDoesNotThrow(() -> hookDispatcher.renameTable(mockContext, request));

    verify(mockDispatcher).renameTable(mockContext, request);
  }

  @Test
  public void testRenameTableThrowsRuntimeExceptionOnIOException() throws IOException {
    TableIdentifier source = TableIdentifier.of("schema1", "old_table");
    TableIdentifier dest = TableIdentifier.of("schema2", "new_table");
    RenameTableRequest request =
        RenameTableRequest.builder().withSource(source).withDestination(dest).build();

    doThrow(new IOException("IO error")).when(mockEntityStore).update(any(), any(), any(), any());

    RuntimeException exception =
        Assertions.assertThrows(
            RuntimeException.class, () -> hookDispatcher.renameTable(mockContext, request));

    Assertions.assertTrue(exception.getMessage().contains("io exception when renaming table"));
    verify(mockDispatcher).renameTable(mockContext, request);
  }

  @Test
  public void testCreateTableThrowsWhenSetOwnerFails() {
    Namespace namespace = Namespace.of("test_schema");
    CreateTableRequest request =
        CreateTableRequest.builder().withName("test_table").withSchema(TABLE_SCHEMA).build();

    LoadTableResponse mockResponse = mock(LoadTableResponse.class);
    when(mockDispatcher.createTable(mockContext, namespace, request)).thenReturn(mockResponse);

    doThrow(new RuntimeException("Set owner failed"))
        .when(mockInternalOwnerDispatcher)
        .setOwner(any(), any(), any(), any());

    RuntimeException thrown =
        Assertions.assertThrows(
            RuntimeException.class,
            () -> hookDispatcher.createTable(mockContext, namespace, request));
    Assertions.assertEquals("Set owner failed", thrown.getMessage());
    verify(mockDispatcher).createTable(mockContext, namespace, request);
    verify(mockOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
  }

  @Test
  public void testCreateTableSkipsImportAndOwnershipForStageCreate() {
    Namespace namespace = Namespace.of("test_schema");
    CreateTableRequest request =
        CreateTableRequest.builder()
            .withName("test_table")
            .withSchema(TABLE_SCHEMA)
            .stageCreate()
            .build();

    LoadTableResponse mockResponse = mock(LoadTableResponse.class);
    when(mockDispatcher.createTable(mockContext, namespace, request)).thenReturn(mockResponse);

    LoadTableResponse result = hookDispatcher.createTable(mockContext, namespace, request);

    Assertions.assertEquals(mockResponse, result);
    verify(mockDispatcher).createTable(mockContext, namespace, request);

    // Verify table import was NOT called for staged create
    verify(mockInternalTableDispatcher, never()).loadTable(any());
    verify(mockTableDispatcher, never()).loadTable(any());

    // Verify ownership was NOT set for staged create
    verify(mockInternalOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
    verify(mockOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
  }

  @Test
  public void testUpdateTableImportsAndSetsOwnershipForStagedCommit() {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");
    // A staged create commit carries AssertTableDoesNotExist as its requirement.
    UpdateTableRequest request =
        new UpdateTableRequest(
            List.of(new UpdateRequirement.AssertTableDoesNotExist()), Collections.emptyList());
    LoadTableResponse mockResponse = mock(LoadTableResponse.class);

    when(mockDispatcher.updateTable(mockContext, tableId, request)).thenReturn(mockResponse);

    LoadTableResponse result = hookDispatcher.updateTable(mockContext, tableId, request);

    Assertions.assertEquals(mockResponse, result);
    verify(mockDispatcher).updateTable(mockContext, tableId, request);

    NameIdentifier gravitinoTableId =
        IcebergIdentifierUtils.toGravitinoTableIdentifier(
            TEST_METALAKE, TEST_CATALOG, tableId, ":");
    verify(mockInternalTableDispatcher).loadTable(gravitinoTableId);
    verify(mockTableDispatcher, never()).loadTable(any());

    ArgumentCaptor<String> userCaptor = ArgumentCaptor.forClass(String.class);
    verify(mockInternalOwnerDispatcher)
        .setOwner(eq(TEST_METALAKE), any(), userCaptor.capture(), eq(Owner.Type.USER));
    verify(mockOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
    Assertions.assertEquals(TEST_USER, userCaptor.getValue());
  }

  @Test
  public void testUpdateTableSkipsImportForRegularUpdate() {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");
    // Regular table update (e.g. property change): no AssertTableDoesNotExist requirement.
    // Import and ownership must NOT be triggered regardless of entity store state.
    UpdateTableRequest request =
        new UpdateTableRequest(Collections.emptyList(), Collections.emptyList());
    LoadTableResponse mockResponse = mock(LoadTableResponse.class);

    when(mockDispatcher.updateTable(mockContext, tableId, request)).thenReturn(mockResponse);

    LoadTableResponse result = hookDispatcher.updateTable(mockContext, tableId, request);

    Assertions.assertEquals(mockResponse, result);
    verify(mockDispatcher).updateTable(mockContext, tableId, request);

    // Verify table import was NOT called for a regular update
    verify(mockInternalTableDispatcher, never()).loadTable(any());
    verify(mockTableDispatcher, never()).loadTable(any());

    // Verify ownership was NOT set
    verify(mockInternalOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
    verify(mockOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
  }

  @Test
  public void testLoadTablePassesThrough() {
    TableIdentifier tableId = TableIdentifier.of("test_schema", "test_table");
    LoadTableResponse mockResponse = mock(LoadTableResponse.class);

    when(mockDispatcher.loadTable(mockContext, tableId)).thenReturn(mockResponse);

    LoadTableResponse result = hookDispatcher.loadTable(mockContext, tableId);

    Assertions.assertEquals(mockResponse, result);
    verify(mockDispatcher).loadTable(mockContext, tableId);
  }

  @Test
  public void testCreateTablePropagatesImportFailure() {
    Namespace namespace = Namespace.of("test_schema");
    CreateTableRequest request =
        CreateTableRequest.builder().withName("test_table").withSchema(TABLE_SCHEMA).build();

    LoadTableResponse mockResponse = mock(LoadTableResponse.class);
    when(mockDispatcher.createTable(mockContext, namespace, request)).thenReturn(mockResponse);

    // Import failure (the loadTable call) must propagate so the caller learns the table exists in
    // Iceberg but is not registered in Gravitino. setOwner is therefore unreachable.
    doThrow(new RuntimeException("Import failed"))
        .when(mockInternalTableDispatcher)
        .loadTable(any());

    RuntimeException thrown =
        Assertions.assertThrows(
            RuntimeException.class,
            () -> hookDispatcher.createTable(mockContext, namespace, request));

    Assertions.assertEquals("Import failed", thrown.getMessage());
    verify(mockDispatcher).createTable(mockContext, namespace, request);
    verify(mockInternalTableDispatcher).loadTable(any());
    verify(mockTableDispatcher, never()).loadTable(any());
    verify(mockInternalOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
    verify(mockOwnerDispatcher, never()).setOwner(any(), any(), any(), any());
  }
}
