]> git.openstreetmap.org Git - nominatim.git/blob - test/python/mock_icu_word_table.py
allow multiple files for the import command
[nominatim.git] / test / python / mock_icu_word_table.py
1 """
2 Legacy word table for testing with functions to prefil and test contents
3 of the table.
4 """
5
6 class MockIcuWordTable:
7     """ A word table for testing using legacy word table structure.
8     """
9     def __init__(self, conn):
10         self.conn = conn
11         with conn.cursor() as cur:
12             cur.execute("""CREATE TABLE word (word_id INTEGER,
13                                               word_token text NOT NULL,
14                                               type text NOT NULL,
15                                               word text,
16                                               info jsonb)""")
17
18         conn.commit()
19
20     def add_special(self, word_token, word, cls, typ, oper):
21         with self.conn.cursor() as cur:
22             cur.execute("""INSERT INTO word (word_token, type, word, info)
23                               VALUES (%s, 'S', %s,
24                                       json_build_object('class', %s,
25                                                         'type', %s,
26                                                         'op', %s))
27                         """, (word_token, word, cls, typ, oper))
28         self.conn.commit()
29
30
31     def add_country(self, country_code, word_token):
32         with self.conn.cursor() as cur:
33             cur.execute("""INSERT INTO word (word_token, type, word)
34                            VALUES(%s, 'C', %s)""",
35                         (word_token, country_code))
36         self.conn.commit()
37
38
39     def add_postcode(self, word_token, postcode):
40         with self.conn.cursor() as cur:
41             cur.execute("""INSERT INTO word (word_token, type, word)
42                               VALUES (%s, 'P', %s)
43                         """, (word_token, postcode))
44         self.conn.commit()
45
46
47     def count(self):
48         with self.conn.cursor() as cur:
49             return cur.scalar("SELECT count(*) FROM word")
50
51
52     def count_special(self):
53         with self.conn.cursor() as cur:
54             return cur.scalar("SELECT count(*) FROM word WHERE type = 'S'")
55
56
57     def get_special(self):
58         with self.conn.cursor() as cur:
59             cur.execute("SELECT word_token, info, word FROM word WHERE type = 'S'")
60             result = set(((row[0], row[2], row[1]['class'],
61                            row[1]['type'], row[1]['op']) for row in cur))
62             assert len(result) == cur.rowcount, "Word table has duplicates."
63             return result
64
65
66     def get_country(self):
67         with self.conn.cursor() as cur:
68             cur.execute("SELECT word, word_token FROM word WHERE type = 'C'")
69             result = set((tuple(row) for row in cur))
70             assert len(result) == cur.rowcount, "Word table has duplicates."
71             return result
72
73
74     def get_postcodes(self):
75         with self.conn.cursor() as cur:
76             cur.execute("SELECT word FROM word WHERE type = 'P'")
77             return set((row[0] for row in cur))
78
79
80     def get_partial_words(self):
81         with self.conn.cursor() as cur:
82             cur.execute("SELECT word_token, info FROM word WHERE type ='w'")
83             return set(((row[0], row[1]['count']) for row in cur))
84