|
| 1 | +"""Tests for the TenThousandFunction.""" |
| 2 | + |
| 3 | +import pyarrow as pa |
| 4 | + |
| 5 | +from vgi.examples.table import TenThousandFunction |
| 6 | + |
| 7 | +from .conftest import RunnerWithMode, run_in_process |
| 8 | + |
| 9 | + |
| 10 | +class TestTenThousandFunctionInProcess: |
| 11 | + """In-process tests for the ten_thousand function.""" |
| 12 | + |
| 13 | + def test_metadata(self) -> None: |
| 14 | + """TenThousand function should have correct metadata.""" |
| 15 | + meta = TenThousandFunction.get_metadata() |
| 16 | + assert meta.name == "ten_thousand" |
| 17 | + assert meta.max_workers == 1 |
| 18 | + assert "generator" in meta.categories |
| 19 | + |
| 20 | + def test_output_schema(self) -> None: |
| 21 | + """Output schema should have single int64 column named 'n'.""" |
| 22 | + outputs, _ = run_in_process(TenThousandFunction, ()) |
| 23 | + assert len(outputs) > 0 |
| 24 | + schema = outputs[0].schema |
| 25 | + assert schema.names == ["n"] |
| 26 | + assert schema.field("n").type == pa.int64() |
| 27 | + |
| 28 | + |
| 29 | +class TestTenThousandFunctionBothModes: |
| 30 | + """Tests that run both in-process and via Client subprocess.""" |
| 31 | + |
| 32 | + def test_generates_ten_thousand_rows( |
| 33 | + self, run_table_function_mode: RunnerWithMode |
| 34 | + ) -> None: |
| 35 | + """TenThousand should generate exactly 10000 rows.""" |
| 36 | + runner, mode = run_table_function_mode |
| 37 | + outputs, logs = runner(TenThousandFunction, ()) |
| 38 | + |
| 39 | + table = pa.Table.from_batches(outputs) |
| 40 | + assert table.num_rows == 10000 |
| 41 | + |
| 42 | + def test_generates_correct_values( |
| 43 | + self, run_table_function_mode: RunnerWithMode |
| 44 | + ) -> None: |
| 45 | + """TenThousand should generate integers from 0 to 9999.""" |
| 46 | + runner, mode = run_table_function_mode |
| 47 | + outputs, logs = runner(TenThousandFunction, ()) |
| 48 | + |
| 49 | + table = pa.Table.from_batches(outputs) |
| 50 | + values = table.column("n").to_pylist() |
| 51 | + assert values == list(range(10000)) |
| 52 | + |
| 53 | + def test_batching(self, run_table_function_mode: RunnerWithMode) -> None: |
| 54 | + """TenThousand should produce batches of 1000 rows each.""" |
| 55 | + runner, mode = run_table_function_mode |
| 56 | + outputs, logs = runner(TenThousandFunction, ()) |
| 57 | + |
| 58 | + # Should produce 10 batches of 1000 rows each |
| 59 | + assert len(outputs) == 10 |
| 60 | + for batch in outputs: |
| 61 | + assert batch.num_rows == 1000 |
0 commit comments