test/python/mocks.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) 2024 by the Nominatim developer community.
   6 # For a full list of authors see the git log.
   7 """
   8 Custom mocks for testing.
   9 """
  10 import itertools
  11
  12 import psycopg2.extras
  13
  14 from nominatim_db.db import properties
  15
  16 # This must always point to the mock word table for the default tokenizer.
  17 from mock_icu_word_table import MockIcuWordTable as MockWordTable
  18
  19 class MockPlacexTable:
  20     """ A placex table for testing.
  21     """
  22     def __init__(self, conn):
  23         self.idseq = itertools.count(10000)
  24         self.conn = conn
  25         with conn.cursor() as cur:
  26             cur.execute("""CREATE TABLE placex (
  27                                place_id BIGINT,
  28                                parent_place_id BIGINT,
  29                                linked_place_id BIGINT,
  30                                importance FLOAT,
  31                                indexed_date TIMESTAMP,
  32                                geometry_sector INTEGER,
  33                                rank_address SMALLINT,
  34                                rank_search SMALLINT,
  35                                partition SMALLINT,
  36                                indexed_status SMALLINT,
  37                                osm_id int8,
  38                                osm_type char(1),
  39                                class text,
  40                                type text,
  41                                name hstore,
  42                                admin_level smallint,
  43                                address hstore,
  44                                extratags hstore,
  45                                token_info jsonb,
  46                                geometry Geometry(Geometry,4326),
  47                                wikipedia TEXT,
  48                                country_code varchar(2),
  49                                housenumber TEXT,
  50                                postcode TEXT,
  51                                centroid GEOMETRY(Geometry, 4326))""")
  52             cur.execute("CREATE SEQUENCE IF NOT EXISTS seq_place")
  53         conn.commit()
  54
  55     def add(self, osm_type='N', osm_id=None, cls='amenity', typ='cafe', names=None,
  56             admin_level=None, address=None, extratags=None, geom='POINT(10 4)',
  57             country=None, housenumber=None, rank_search=30):
  58         with self.conn.cursor() as cur:
  59             psycopg2.extras.register_hstore(cur)
  60             cur.execute("""INSERT INTO placex (place_id, osm_type, osm_id, class,
  61                                                type, name, admin_level, address,
  62                                                housenumber, rank_search,
  63                                                extratags, geometry, country_code)
  64                             VALUES(nextval('seq_place'), %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""",
  65                         (osm_type, osm_id or next(self.idseq), cls, typ, names,
  66                          admin_level, address, housenumber, rank_search,
  67                          extratags, 'SRID=4326;' + geom,
  68                          country))
  69         self.conn.commit()
  70
  71
  72 class MockPropertyTable:
  73     """ A property table for testing.
  74     """
  75     def __init__(self, conn):
  76         self.conn = conn
  77
  78
  79     def set(self, name, value):
  80         """ Set a property in the table to the given value.
  81         """
  82         properties.set_property(self.conn, name, value)
  83
  84
  85     def get(self, name):
  86         """ Set a property in the table to the given value.
  87         """
  88         return properties.get_property(self.conn, name)