test/python/tools/test_tiger_data.py

   1 """
   2 Test for tiger data function
   3 """
   4 import tarfile
   5 from textwrap import dedent
   6
   7 import pytest
   8
   9 from nominatim.tools import tiger_data
  10 from nominatim.errors import UsageError
  11
  12 class MockTigerTable:
  13
  14     def __init__(self, conn):
  15         self.conn = conn
  16         with conn.cursor() as cur:
  17             cur.execute("""CREATE TABLE tiger (linegeo GEOMETRY,
  18                                                start INTEGER,
  19                                                stop INTEGER,
  20                                                interpol TEXT,
  21                                                token_info JSONB,
  22                                                postcode TEXT)""")
  23
  24     def count(self):
  25         with self.conn.cursor() as cur:
  26             return cur.scalar("SELECT count(*) FROM tiger")
  27
  28     def row(self):
  29         with self.conn.cursor() as cur:
  30             cur.execute("SELECT * FROM tiger LIMIT 1")
  31             return cur.fetchone()
  32
  33 @pytest.fixture
  34 def tiger_table(def_config, temp_db_conn, sql_preprocessor,
  35                 temp_db_with_extensions, tmp_path):
  36     def_config.lib_dir.sql = tmp_path / 'sql'
  37     def_config.lib_dir.sql.mkdir()
  38
  39     (def_config.lib_dir.sql / 'tiger_import_start.sql').write_text(
  40         """CREATE OR REPLACE FUNCTION tiger_line_import(linegeo GEOMETRY, start INTEGER,
  41                                                         stop INTEGER, interpol TEXT,
  42                                                         token_info JSONB, postcode TEXT)
  43            RETURNS INTEGER AS $$
  44             INSERT INTO tiger VALUES(linegeo, start, stop, interpol, token_info, postcode)
  45             RETURNING 1
  46            $$ LANGUAGE SQL;""")
  47     (def_config.lib_dir.sql / 'tiger_import_finish.sql').write_text(
  48         """DROP FUNCTION tiger_line_import (linegeo GEOMETRY, in_startnumber INTEGER,
  49                                  in_endnumber INTEGER, interpolationtype TEXT,
  50                                  token_info JSONB, in_postcode TEXT);""")
  51
  52     return MockTigerTable(temp_db_conn)
  53
  54
  55 @pytest.fixture
  56 def csv_factory(tmp_path):
  57     def _mk_file(fname, hnr_from=1, hnr_to=9, interpol='odd', street='Main St',
  58                  city='Newtown', state='AL', postcode='12345',
  59                  geometry='LINESTRING(-86.466995 32.428956,-86.466923 32.428933)'):
  60         (tmp_path / (fname + '.csv')).write_text(dedent("""\
  61         from;to;interpolation;street;city;state;postcode;geometry
  62         {};{};{};{};{};{};{};{}
  63         """.format(hnr_from, hnr_to, interpol, street, city, state,
  64                    postcode, geometry)))
  65
  66     return _mk_file
  67
  68
  69 @pytest.mark.parametrize("threads", (1, 5))
  70 def test_add_tiger_data(def_config, src_dir, tiger_table, tokenizer_mock, threads):
  71     tiger_data.add_tiger_data(str(src_dir / 'test' / 'testdb' / 'tiger'),
  72                               def_config, threads, tokenizer_mock())
  73
  74     assert tiger_table.count() == 6213
  75
  76
  77 def test_add_tiger_data_no_files(def_config, tiger_table, tokenizer_mock,
  78                                  tmp_path):
  79     tiger_data.add_tiger_data(str(tmp_path), def_config, 1, tokenizer_mock())
  80
  81     assert tiger_table.count() == 0
  82
  83
  84 def test_add_tiger_data_bad_file(def_config, tiger_table, tokenizer_mock,
  85                                  tmp_path):
  86     sqlfile = tmp_path / '1010.csv'
  87     sqlfile.write_text("""Random text""")
  88
  89     tiger_data.add_tiger_data(str(tmp_path), def_config, 1, tokenizer_mock())
  90
  91     assert tiger_table.count() == 0
  92
  93
  94 def test_add_tiger_data_hnr_nan(def_config, tiger_table, tokenizer_mock,
  95                                 csv_factory, tmp_path):
  96     csv_factory('file1', hnr_from=99)
  97     csv_factory('file2', hnr_from='L12')
  98     csv_factory('file3', hnr_to='12.4')
  99
 100     tiger_data.add_tiger_data(str(tmp_path), def_config, 1, tokenizer_mock())
 101
 102     assert tiger_table.count() == 1
 103     assert tiger_table.row()['start'] == 99
 104
 105
 106 @pytest.mark.parametrize("threads", (1, 5))
 107 def test_add_tiger_data_tarfile(def_config, tiger_table, tokenizer_mock,
 108                                 tmp_path, src_dir, threads):
 109     tar = tarfile.open(str(tmp_path / 'sample.tar.gz'), "w:gz")
 110     tar.add(str(src_dir / 'test' / 'testdb' / 'tiger' / '01001.csv'))
 111     tar.close()
 112
 113     tiger_data.add_tiger_data(str(tmp_path / 'sample.tar.gz'), def_config, threads,
 114                               tokenizer_mock())
 115
 116     assert tiger_table.count() == 6213
 117
 118
 119 def test_add_tiger_data_bad_tarfile(def_config, tiger_table, tokenizer_mock,
 120                                     tmp_path):
 121     tarfile = tmp_path / 'sample.tar.gz'
 122     tarfile.write_text("""Random text""")
 123
 124     with pytest.raises(UsageError):
 125         tiger_data.add_tiger_data(str(tarfile), def_config, 1, tokenizer_mock())
 126
 127
 128 def test_add_tiger_data_empty_tarfile(def_config, tiger_table, tokenizer_mock,
 129                                       tmp_path):
 130     tar = tarfile.open(str(tmp_path / 'sample.tar.gz'), "w:gz")
 131     tar.add(__file__)
 132     tar.close()
 133
 134     tiger_data.add_tiger_data(str(tmp_path / 'sample.tar.gz'), def_config, 1,
 135                               tokenizer_mock())
 136
 137     assert tiger_table.count() == 0