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