test/python/tools/test_tiger_data.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2025 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 import pytest_asyncio  # noqa: F401
  15
  16 from nominatim_db.db.connection import execute_scalar
  17 from nominatim_db.tools import tiger_data, freeze
  18 from nominatim_db.errors import UsageError
  19
  20
  21 class MockTigerTable:
  22
  23     def __init__(self, conn):
  24         self.conn = conn
  25         with conn.cursor() as cur:
  26             cur.execute("""CREATE TABLE tiger (linegeo GEOMETRY,
  27                                                start INTEGER,
  28                                                stop INTEGER,
  29                                                interpol TEXT,
  30                                                token_info JSONB,
  31                                                postcode TEXT)""")
  32
  33             # We need this table to determine if the database is frozen or not
  34             cur.execute("CREATE TABLE place (number INTEGER)")
  35             # We need this table to determine if the database is in reverse-only mode
  36             cur.execute("CREATE TABLE search_name (place_id BIGINT)")
  37
  38     def count(self):
  39         return execute_scalar(self.conn, "SELECT count(*) FROM tiger")
  40
  41     def row(self):
  42         with self.conn.cursor() as cur:
  43             cur.execute("SELECT * FROM tiger LIMIT 1")
  44             return cur.fetchone()
  45
  46
  47 @pytest.fixture
  48 def tiger_table(def_config, temp_db_conn, sql_preprocessor,
  49                 temp_db_with_extensions, tmp_path):
  50     def_config.lib_dir.sql = tmp_path / 'sql'
  51     def_config.lib_dir.sql.mkdir()
  52
  53     (def_config.lib_dir.sql / 'tiger_import_start.sql').write_text(
  54         """CREATE OR REPLACE FUNCTION tiger_line_import(linegeo GEOMETRY, start INTEGER,
  55                                                         stop INTEGER, interpol TEXT,
  56                                                         token_info JSONB, postcode TEXT)
  57            RETURNS INTEGER AS $$
  58             INSERT INTO tiger VALUES(linegeo, start, stop, interpol, token_info, postcode)
  59             RETURNING 1
  60            $$ LANGUAGE SQL;""")
  61     (def_config.lib_dir.sql / 'tiger_import_finish.sql').write_text(
  62         """DROP FUNCTION tiger_line_import (linegeo GEOMETRY, in_startnumber INTEGER,
  63                                  in_endnumber INTEGER, interpolationtype TEXT,
  64                                  token_info JSONB, in_postcode TEXT);""")
  65
  66     return MockTigerTable(temp_db_conn)
  67
  68
  69 @pytest.fixture
  70 def csv_factory(tmp_path):
  71     def _mk_file(fname, hnr_from=1, hnr_to=9, interpol='odd', street='Main St',
  72                  city='Newtown', state='AL', postcode='12345',
  73                  geometry='LINESTRING(-86.466995 32.428956,-86.466923 32.428933)'):
  74         (tmp_path / (fname + '.csv')).write_text(dedent("""\
  75         from;to;interpolation;street;city;state;postcode;geometry
  76         {};{};{};{};{};{};{};{}
  77         """.format(hnr_from, hnr_to, interpol, street, city, state,
  78                    postcode, geometry)))
  79
  80     return _mk_file
  81
  82
  83 @pytest.mark.parametrize("threads", (1, 5))
  84 @pytest.mark.asyncio
  85 async def test_add_tiger_data(def_config, src_dir, tiger_table, tokenizer_mock, threads):
  86     await tiger_data.add_tiger_data(str(src_dir / 'test' / 'testdb' / 'tiger'),
  87                                     def_config, threads, tokenizer_mock())
  88
  89     assert tiger_table.count() == 6213
  90
  91
  92 @pytest.mark.parametrize("threads", (1, 5))
  93 @pytest.mark.asyncio
  94 async def test_add_tiger_data_database_frozen(def_config, src_dir, temp_db_conn, tiger_table,
  95                                               tokenizer_mock, threads):
  96     freeze.drop_update_tables(temp_db_conn)
  97
  98     await tiger_data.add_tiger_data(str(src_dir / 'test' / 'testdb' / 'tiger'),
  99                                     def_config, threads, tokenizer_mock())
 100
 101     assert tiger_table.count() == 6213
 102
 103
 104 @pytest.mark.asyncio
 105 async def test_add_tiger_data_reverse_only(def_config, src_dir, temp_db_conn, tiger_table,
 106                                            tokenizer_mock):
 107     with temp_db_conn.cursor() as cur:
 108         cur.execute("DROP TABLE search_name")
 109     temp_db_conn.commit()
 110
 111     with pytest.raises(UsageError,
 112                        match="Cannot perform tiger import: required tables are missing. "
 113                        "See https://github.com/osm-search/Nominatim/issues/2463 for details."):
 114         await tiger_data.add_tiger_data(str(src_dir / 'test' / 'testdb' / 'tiger'),
 115                                         def_config, 1, tokenizer_mock())
 116
 117     assert tiger_table.count() == 0
 118
 119
 120 @pytest.mark.asyncio
 121 async def test_add_tiger_data_no_files(def_config, tiger_table, tokenizer_mock,
 122                                        tmp_path):
 123     await tiger_data.add_tiger_data(str(tmp_path), def_config, 1, tokenizer_mock())
 124
 125     assert tiger_table.count() == 0
 126
 127
 128 @pytest.mark.asyncio
 129 async def test_add_tiger_data_bad_file(def_config, tiger_table, tokenizer_mock,
 130                                        tmp_path):
 131     sqlfile = tmp_path / '1010.csv'
 132     sqlfile.write_text("""Random text""")
 133
 134     await tiger_data.add_tiger_data(str(tmp_path), def_config, 1, tokenizer_mock())
 135
 136     assert tiger_table.count() == 0
 137
 138
 139 @pytest.mark.asyncio
 140 async def test_add_tiger_data_hnr_nan(def_config, tiger_table, tokenizer_mock,
 141                                       csv_factory, tmp_path):
 142     csv_factory('file1', hnr_from=99)
 143     csv_factory('file2', hnr_from='L12')
 144     csv_factory('file3', hnr_to='12.4')
 145
 146     await tiger_data.add_tiger_data(str(tmp_path), def_config, 1, tokenizer_mock())
 147
 148     assert tiger_table.count() == 1
 149     assert tiger_table.row().start == 99
 150
 151
 152 @pytest.mark.parametrize("threads", (1, 5))
 153 @pytest.mark.asyncio
 154 async def test_add_tiger_data_tarfile(def_config, tiger_table, tokenizer_mock,
 155                                       tmp_path, src_dir, threads):
 156     tar = tarfile.open(str(tmp_path / 'sample.tar.gz'), "w:gz")
 157     tar.add(str(src_dir / 'test' / 'testdb' / 'tiger' / '01001.csv'))
 158     tar.close()
 159
 160     await tiger_data.add_tiger_data(str(tmp_path / 'sample.tar.gz'), def_config, threads,
 161                                     tokenizer_mock())
 162
 163     assert tiger_table.count() == 6213
 164
 165
 166 @pytest.mark.asyncio
 167 async def test_add_tiger_data_bad_tarfile(def_config, tiger_table, tokenizer_mock,
 168                                           tmp_path):
 169     tarfile = tmp_path / 'sample.tar.gz'
 170     tarfile.write_text("""Random text""")
 171
 172     with pytest.raises(UsageError):
 173         await tiger_data.add_tiger_data(str(tarfile), def_config, 1, tokenizer_mock())
 174
 175
 176 @pytest.mark.asyncio
 177 async def test_add_tiger_data_empty_tarfile(def_config, tiger_table, tokenizer_mock,
 178                                             tmp_path):
 179     tar = tarfile.open(str(tmp_path / 'sample.tar.gz'), "w:gz")
 180     tar.add(__file__)
 181     tar.close()
 182
 183     await tiger_data.add_tiger_data(str(tmp_path / 'sample.tar.gz'), def_config, 1,
 184                                     tokenizer_mock())
 185
 186     assert tiger_table.count() == 0