wrds-download

TUI/CLI tool for browsing and downloading WRDS data
Log | Files | Refs | README

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)) == ["", ""]