Add a test

huggingface · Nov 8, 2021 · e8f6fae · e8f6fae · github-actions · Nov 8, 2021
1 parent daaa0de
commit e8f6fae
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/tests/test_arrow_writer.py b/tests/test_arrow_writer.py
@@ -234,6 +234,18 @@ def test_optimized_typed_sequence(sequence, col, expected_dtype):
     assert get_base_dtype(arr.type) == expected_dtype
 
 
+def test_arrow_writer_typed_sequence_cls(monkeypatch):
+    stream = pa.BufferOutputStream()
+
+    with ArrowWriter(stream=stream) as writer:
+        assert writer.typed_sequence_cls == OptimizedTypedSequence
+
+    monkeypatch.setattr(config, "OPTIMIZE_PYARROW_TYPES", False)
+
+    with ArrowWriter(stream=stream) as writer:
+        assert writer.typed_sequence_cls == TypedSequence
+
+
 @pytest.mark.parametrize("raise_exception", [False, True])
 def test_arrow_writer_closes_stream(raise_exception, tmp_path):
     path = str(tmp_path / "dataset-train.arrow")