test/python/tools/test_tiger_data.py

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