iceberg/
test_utils.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Test utilities.
19//! This module is pub just for internal testing.
20//! It is subject to change and is not intended to be used by external users.
21
22use std::sync::{Arc, OnceLock};
23
24use arrow_array::RecordBatch;
25use expect_test::Expect;
26use itertools::Itertools;
27#[cfg(test)]
28use roaring::RoaringTreemap;
29
30use crate::TableIdent;
31#[cfg(test)]
32use crate::encryption::EncryptionManager;
33use crate::encryption::SensitiveBytes;
34use crate::encryption::kms::{KeyManagementClient, MemoryKeyManagementClient};
35use crate::io::FileIO;
36use crate::runtime::Runtime;
37use crate::spec::TableMetadata;
38use crate::table::Table;
39
40/// Returns a process-wide [`Runtime`] suitable for tests that need to construct
41/// a [`Table`] outside a tokio context.
42///
43/// The returned [`Runtime`] wraps a single shared multi-thread tokio runtime
44/// that is lazily built on first call and lives until process exit. Cloning is
45/// cheap, so test code can call this every time it needs a runtime to feed
46/// into [`TableBuilder::runtime`](crate::table::TableBuilder::runtime).
47pub fn test_runtime() -> Runtime {
48    static TOKIO_RT: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
49    let tokio_rt = TOKIO_RT.get_or_init(|| {
50        tokio::runtime::Builder::new_multi_thread()
51            .enable_all()
52            .build()
53            .expect("failed to build test tokio runtime")
54    });
55    Runtime::new(tokio_rt)
56}
57
58/// Snapshot testing to check the resulting record batch.
59///
60/// - `expected_schema/data`: put `expect![[""]]` as a placeholder,
61///   and then run test with `UPDATE_EXPECT=1 cargo test` to automatically update the result,
62///   or use rust-analyzer (see [video](https://github.com/rust-analyzer/expect-test)).
63///   Check the doc of [`expect_test`] for more details.
64/// - `ignore_check_columns`: Some columns are not stable, so we can skip them.
65/// - `sort_column`: The order of the data might be non-deterministic, so we can sort it by a column.
66pub fn check_record_batches(
67    record_batches: Vec<RecordBatch>,
68    expected_schema: Expect,
69    expected_data: Expect,
70    ignore_check_columns: &[&str],
71    sort_column: Option<&str>,
72) {
73    assert!(!record_batches.is_empty(), "Empty record batches");
74
75    // Combine record batches using the first batch's schema
76    let first_batch = record_batches.first().unwrap();
77    let record_batch =
78        arrow_select::concat::concat_batches(&first_batch.schema(), &record_batches).unwrap();
79
80    let mut columns = record_batch.columns().to_vec();
81    if let Some(sort_column) = sort_column {
82        let column = record_batch.column_by_name(sort_column).unwrap();
83        let indices = arrow_ord::sort::sort_to_indices(column, None, None).unwrap();
84        columns = columns
85            .iter()
86            .map(|column| arrow_select::take::take(column.as_ref(), &indices, None).unwrap())
87            .collect_vec();
88    }
89
90    expected_schema.assert_eq(&format!(
91        "{}",
92        record_batch.schema().fields().iter().format(",\n")
93    ));
94    expected_data.assert_eq(&format!(
95        "{}",
96        record_batch
97            .schema()
98            .fields()
99            .iter()
100            .zip_eq(columns)
101            .map(|(field, column)| {
102                if ignore_check_columns.contains(&field.name().as_str()) {
103                    format!("{}: (skipped)", field.name())
104                } else {
105                    format!("{}: {:?}", field.name(), column)
106                }
107            })
108            .format(",\n")
109    ));
110}
111
112/// An [`EncryptionManager`] backed by an in-memory KMS holding `table_key_id`.
113#[cfg(test)]
114pub(crate) fn make_encryption_manager(table_key_id: &str) -> Arc<EncryptionManager> {
115    let kms = MemoryKeyManagementClient::new();
116    kms.add_master_key(table_key_id).unwrap();
117    Arc::new(
118        EncryptionManager::builder()
119            .kms_client(Arc::new(kms) as Arc<dyn KeyManagementClient>)
120            .table_key_id(table_key_id)
121            .build(),
122    )
123}
124
125/// Encodes a `deletion-vector-v1` Puffin blob for the given positions, matching the framing in
126/// [`DeleteVector::deserialize`](crate::delete_vector::DeleteVector::deserialize).
127#[cfg(test)]
128pub(crate) fn encode_dv_blob(positions: impl IntoIterator<Item = u64>) -> Vec<u8> {
129    let mut bitmap = RoaringTreemap::new();
130    for pos in positions {
131        bitmap.insert(pos);
132    }
133    let mut vector = Vec::new();
134    bitmap.serialize_into(&mut vector).unwrap();
135    crate::delete_vector::frame_dv_blob(&vector)
136}
137
138/// Build a table backed by the V3 encryption fixture and an in-memory KMS,
139/// so it has an [`EncryptionManager`](crate::encryption::EncryptionManager).
140///
141/// The fixture's snapshot references an encrypted manifest list; its bytes
142/// (the `manifest-list-v3-encrypted.avro` testdata, an encrypted empty list)
143/// are seeded into the in-memory `FileIO` at that path so callers can read
144/// the current snapshot's manifest list.
145pub async fn make_encrypted_table() -> Table {
146    let metadata_json = std::fs::read_to_string(format!(
147        "{}/testdata/table_metadata/TableMetadataV3ValidEncryption.json",
148        env!("CARGO_MANIFEST_DIR"),
149    ))
150    .unwrap();
151    let metadata: TableMetadata = serde_json::from_str(&metadata_json).unwrap();
152
153    let kms: Arc<dyn KeyManagementClient> = {
154        let k = MemoryKeyManagementClient::new();
155        k.add_master_key_bytes(
156            "master-1",
157            SensitiveBytes::new([
158                0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
159                0x0e, 0x0f,
160            ]),
161        )
162        .unwrap();
163        Arc::new(k)
164    };
165
166    let file_io = FileIO::new_with_memory();
167
168    // Seed the encrypted (empty) manifest list at the path the snapshot references.
169    let manifest_list_bytes = std::fs::read(format!(
170        "{}/testdata/manifests_lists/manifest-list-v3-encrypted.avro",
171        env!("CARGO_MANIFEST_DIR"),
172    ))
173    .unwrap();
174    file_io
175        .new_output(metadata.current_snapshot().unwrap().manifest_list())
176        .unwrap()
177        .write(manifest_list_bytes.into())
178        .await
179        .unwrap();
180
181    Table::builder()
182        .metadata(metadata)
183        .metadata_location("memory:///table/metadata/v1.json")
184        .identifier(TableIdent::from_strs(["ns1", "test1"]).unwrap())
185        .file_io(file_io)
186        .kms_client(kms)
187        .runtime(test_runtime())
188        .build()
189        .unwrap()
190}
iceberg/test_utils.rs

iceberg/
test_utils.rs
