src/iceberg/test/update_test_base.h

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

#pragma once

#include <format>
#include <functional>
#include <memory>
#include <string>

#include <arrow/filesystem/mockfs.h>
#include <gtest/gtest.h>

#include "iceberg/arrow/arrow_io_internal.h"
#include "iceberg/catalog/memory/in_memory_catalog.h"
#include "iceberg/result.h"
#include "iceberg/snapshot.h"
#include "iceberg/table.h"
#include "iceberg/table_identifier.h"
#include "iceberg/table_metadata.h"
#include "iceberg/test/matchers.h"
#include "iceberg/test/mock_catalog.h"
#include "iceberg/test/test_resource.h"
#include "iceberg/util/uuid.h"

namespace iceberg {

/// \brief Base test fixture for table update operations.
class UpdateTestBase : public ::testing::Test {
 protected:
  virtual std::string MetadataResource() const { return "TableMetadataV2Valid.json"; }
  virtual std::string TableName() const { return "test_table"; }

  void SetUp() override {
    table_ident_ = TableIdentifier{.name = TableName()};
    table_location_ = "/warehouse/" + TableName();

    InitializeFileIO();
    RegisterTableFromResource(MetadataResource());
  }

  /// \brief Initialize file IO and create necessary directories.
  void InitializeFileIO() {
    file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO();
    ICEBERG_UNWRAP_OR_FAIL(catalog_,
                           InMemoryCatalog::Make("test_catalog", file_io_, "/warehouse/",
                                                 /*properties=*/{}));

    // Arrow MockFS cannot automatically create directories.
    auto arrow_fs = std::dynamic_pointer_cast<::arrow::fs::internal::MockFileSystem>(
        static_cast<arrow::ArrowFileSystemFileIO&>(*file_io_).fs());
    ASSERT_TRUE(arrow_fs != nullptr);
    ASSERT_TRUE(arrow_fs->CreateDir(table_location_ + "/metadata").ok());
    ASSERT_TRUE(arrow_fs->CreateDir(table_location_ + "/data").ok());
  }

  /// \brief Register a table from a metadata resource file.
  ///
  /// \param resource_name The name of the metadata resource file
  void RegisterTableFromResource(const std::string& resource_name) {
    std::ignore = catalog_->DropTable(table_ident_, /*purge=*/false);

    auto metadata_location = std::format("{}/metadata/00001-{}.metadata.json",
                                         table_location_, Uuid::GenerateV7().ToString());
    ICEBERG_UNWRAP_OR_FAIL(auto metadata, ReadTableMetadataFromResource(resource_name));
    metadata->location = table_location_;
    ASSERT_THAT(TableMetadataUtil::Write(*file_io_, metadata_location, *metadata),
                IsOk());

    ICEBERG_UNWRAP_OR_FAIL(table_,
                           catalog_->RegisterTable(table_ident_, metadata_location));
  }

  // Exercise the public commit retry path against real catalog metadata.
  void FailCommits(int conflicts, std::function<void(int)> before_attempt = {}) {
    auto mock = std::make_shared<::testing::NiceMock<MockCatalog>>();
    EXPECT_CALL(*mock, LoadTable(::testing::_))
        .Times(::testing::AtLeast(conflicts))
        .WillRepeatedly([catalog = catalog_](const TableIdentifier& name) {
          return catalog->LoadTable(name);
        });
    EXPECT_CALL(*mock, UpdateTable(::testing::_, ::testing::_, ::testing::_))
        .Times(conflicts + 1)
        .WillRepeatedly(
            [catalog = catalog_, conflicts, before_attempt,
             attempt = std::make_shared<int>(0)](
                const auto& name, const auto& requirements,
                const auto& updates) mutable -> Result<std::shared_ptr<Table>> {
              if (before_attempt) {
                before_attempt(*attempt);
              }
              if ((*attempt)++ < conflicts) {
                return CommitFailed("injected conflict");
              }
              auto result = catalog->UpdateTable(name, requirements, updates);
              if (!result) {
                ADD_FAILURE() << result.error().message;
                return ValidationFailed("Test catalog failed: {}",
                                        result.error().message);
              }
              return result;
            });
    ICEBERG_UNWRAP_OR_FAIL(
        table_, Table::Make(table_->name(), table_->metadata(),
                            std::string(table_->metadata_file_location()), file_io_, mock,
                            table_->full_name(), table_->reporter()));
  }

  /// \brief Reload the table from catalog and return its metadata.
  std::shared_ptr<TableMetadata> ReloadMetadata() {
    auto result = catalog_->LoadTable(table_ident_);
    EXPECT_TRUE(result.has_value()) << "Failed to reload table";
    return result.value()->metadata();
  }

  /// \brief Assert that a ref exists with the given type and snapshot id.
  void ExpectRef(const std::string& name, SnapshotRefType type, int64_t snapshot_id) {
    auto metadata = ReloadMetadata();
    auto it = metadata->refs.find(name);
    ASSERT_NE(it, metadata->refs.end()) << "Ref not found: " << name;
    EXPECT_EQ(it->second->type(), type);
    EXPECT_EQ(it->second->snapshot_id, snapshot_id);
  }

  void ExpectBranch(const std::string& name, int64_t snapshot_id) {
    ExpectRef(name, SnapshotRefType::kBranch, snapshot_id);
  }

  void ExpectTag(const std::string& name, int64_t snapshot_id) {
    ExpectRef(name, SnapshotRefType::kTag, snapshot_id);
  }

  /// \brief Assert that a ref does not exist.
  void ExpectNoRef(const std::string& name) {
    auto metadata = ReloadMetadata();
    EXPECT_FALSE(metadata->refs.contains(name)) << "Ref should not exist: " << name;
  }

  /// \brief Assert the current snapshot id after reloading.
  void ExpectCurrentSnapshot(int64_t snapshot_id) {
    auto result = catalog_->LoadTable(table_ident_);
    ASSERT_TRUE(result.has_value());
    auto snap_result = result.value()->current_snapshot();
    ASSERT_TRUE(snap_result.has_value());
    EXPECT_EQ(snap_result.value()->snapshot_id, snapshot_id);
  }

  /// \brief Assert that a commit succeeded.
  template <typename T>
  void ExpectCommitOk(const T& result) {
    EXPECT_THAT(result, IsOk());
  }

  /// \brief Assert that a commit failed with the given error kind and message substring.
  template <typename T>
  void ExpectCommitError(const T& result, ErrorKind kind, const std::string& message) {
    EXPECT_THAT(result, IsError(kind));
    EXPECT_THAT(result, HasErrorMessage(message));
  }

  TableIdentifier table_ident_;
  std::string table_location_;
  std::shared_ptr<FileIO> file_io_;
  std::shared_ptr<InMemoryCatalog> catalog_;
  std::shared_ptr<Table> table_;
};

/// \brief Test fixture for table update operations on minimal table metadata.
class MinimalUpdateTestBase : public UpdateTestBase {
 protected:
  virtual int8_t format_version() const {
    return TableMetadata::kDefaultTableFormatVersion;
  }

  std::string MetadataResource() const override {
    switch (format_version()) {
      case 1:
        return "TableMetadataV1Valid.json";
      case 2:
        return "TableMetadataV2ValidMinimal.json";
      case 3:
        return "TableMetadataV3ValidMinimal.json";
      default:
        ADD_FAILURE() << "Unsupported format version: "
                      << static_cast<int>(format_version());
        return "TableMetadataV2ValidMinimal.json";
    }
  }

  std::string TableName() const override { return "minimal_table"; }
};

}  // namespace iceberg