test_export.py (1327B)
1 """Tests for export module — type mapping and format helpers.""" 2 3 import pyarrow as pa 4 5 from wrds_dl.export import _arrow_type_for_oid, _format_row 6 7 8 class TestArrowTypeForOid: 9 def test_bool(self): 10 assert _arrow_type_for_oid(16) == pa.bool_() 11 12 def test_int2(self): 13 assert _arrow_type_for_oid(21) == pa.int32() 14 15 def test_int4(self): 16 assert _arrow_type_for_oid(23) == pa.int32() 17 18 def test_int8(self): 19 assert _arrow_type_for_oid(20) == pa.int64() 20 21 def test_float4(self): 22 assert _arrow_type_for_oid(700) == pa.float32() 23 24 def test_float8(self): 25 assert _arrow_type_for_oid(701) == pa.float64() 26 27 def test_date(self): 28 assert _arrow_type_for_oid(1082) == pa.date32() 29 30 def test_timestamp(self): 31 assert _arrow_type_for_oid(1114) == pa.timestamp("us") 32 33 def test_timestamptz(self): 34 assert _arrow_type_for_oid(1184) == pa.timestamp("us", tz="UTC") 35 36 def test_unknown_defaults_to_string(self): 37 assert _arrow_type_for_oid(9999) == pa.string() 38 39 40 class TestFormatRow: 41 def test_basic(self): 42 assert _format_row((1, "hello", None, 3.14)) == ["1", "hello", "", "3.14"] 43 44 def test_empty(self): 45 assert _format_row(()) == [] 46 47 def test_all_none(self): 48 assert _format_row((None, None)) == ["", ""]