diff --git a/src-python/amazon/ionbenchmark/ion_load_dump.py b/src-python/amazon/ionbenchmark/ion_load_dump.py index 92202ee87..300e22192 100644 --- a/src-python/amazon/ionbenchmark/ion_load_dump.py +++ b/src-python/amazon/ionbenchmark/ion_load_dump.py @@ -35,7 +35,7 @@ def load(self, fp): def dumps(self, obj): ion.c_ext = self._c_ext - return ion.dumps(obj, binary=self._binary) + return ion.dumps(obj, binary=self._binary, sequence_as_stream=True) def dump(self, obj, fp): ion.c_ext = self._c_ext diff --git a/tests/test_benchmark_spec.py b/tests/test_benchmark_spec.py index 03b0cb77a..4818b51d8 100644 --- a/tests/test_benchmark_spec.py +++ b/tests/test_benchmark_spec.py @@ -4,6 +4,7 @@ from pathlib import Path import cbor2 +import pytest from amazon.ion import simpleion from amazon.ion.equivalence import ion_equals @@ -139,6 +140,26 @@ def test_write_generate_multiple_top_level_cbor_values(): assert obj_count == load_count +@pytest.mark.parametrize('format_option', ['ion_text', 'ion_binary']) +@pytest.mark.parametrize('values', [[], [1], [1, {'name': 'value'}], [[1, 2]]]) +def test_buffer_write_preserves_top_level_values(format_option, values, tmp_path, monkeypatch): + monkeypatch.setattr(simpleion, 'c_ext', False) + input_file = tmp_path / 'input.ion' + with input_file.open('wb') as fp: + simpleion.dump(values, fp, binary=True, sequence_as_stream=True) + params = {'command': 'write', 'format': format_option, 'input_file': str(input_file), + 'py_c_extension': False} + buffer_fun = _create_test_fun(BenchmarkSpec({**params, 'io_type': 'buffer'})) + output_file = tmp_path / 'output.ion' + _create_test_fun(BenchmarkSpec({**params, 'io_type': 'file'}), custom_file=str(output_file))() + with output_file.open('rb') as fp: + file_values = simpleion.load(fp, single_value=False) + assert ion_equals(file_values, values) + # Benchmark functions are called repeatedly with the same input stream. + for _ in range(2): + assert ion_equals(simpleion.loads(buffer_fun(), single_value=False), file_values) + + def test_get_input_file_size(): real_size = _minimal_spec.get_input_file_size() exp_size = Path(join(_generate_test_path("sample_spec"), _minimal_params['input_file'])).stat().st_size