test/python/mocks.py

   1 """
   2 Custom mocks for testing.
   3 """
   4 import itertools
   5
   6 import psycopg2.extras
   7
   8 from nominatim.db import properties
   9
  10 # This must always point to the mock word table for the default tokenizer.
  11 from mock_legacy_word_table import MockLegacyWordTable as MockWordTable
  12
  13 class MockPlacexTable:
  14     """ A placex table for testing.
  15     """
  16     def __init__(self, conn):
  17         self.idseq = itertools.count(10000)
  18         self.conn = conn
  19         with conn.cursor() as cur:
  20             cur.execute("""CREATE TABLE placex (
  21                                place_id BIGINT,
  22                                parent_place_id BIGINT,
  23                                linked_place_id BIGINT,
  24                                importance FLOAT,
  25                                indexed_date TIMESTAMP,
  26                                geometry_sector INTEGER,
  27                                rank_address SMALLINT,
  28                                rank_search SMALLINT,
  29                                partition SMALLINT,
  30                                indexed_status SMALLINT,
  31                                osm_id int8,
  32                                osm_type char(1),
  33                                class text,
  34                                type text,
  35                                name hstore,
  36                                admin_level smallint,
  37                                address hstore,
  38                                extratags hstore,
  39                                geometry Geometry(Geometry,4326),
  40                                wikipedia TEXT,
  41                                country_code varchar(2),
  42                                housenumber TEXT,
  43                                postcode TEXT,
  44                                centroid GEOMETRY(Geometry, 4326))""")
  45             cur.execute("CREATE SEQUENCE IF NOT EXISTS seq_place")
  46         conn.commit()
  47
  48     def add(self, osm_type='N', osm_id=None, cls='amenity', typ='cafe', names=None,
  49             admin_level=None, address=None, extratags=None, geom='POINT(10 4)',
  50             country=None):
  51         with self.conn.cursor() as cur:
  52             psycopg2.extras.register_hstore(cur)
  53             cur.execute("""INSERT INTO placex (place_id, osm_type, osm_id, class,
  54                                                type, name, admin_level, address,
  55                                                extratags, geometry, country_code)
  56                             VALUES(nextval('seq_place'), %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""",
  57                         (osm_type, osm_id or next(self.idseq), cls, typ, names,
  58                          admin_level, address, extratags, 'SRID=4326;' + geom,
  59                          country))
  60         self.conn.commit()
  61
  62
  63 class MockPropertyTable:
  64     """ A property table for testing.
  65     """
  66     def __init__(self, conn):
  67         self.conn = conn
  68
  69
  70     def set(self, name, value):
  71         """ Set a property in the table to the given value.
  72         """
  73         properties.set_property(self.conn, name, value)