X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/430c316e45c5fcbaf7f8023958a13ad1f31f44b2..b4fec57b6d53f8e8a45c46ff11f13cbcbea1006a:/test/python/mocks.py

diff --git a/test/python/mocks.py b/test/python/mocks.py
index d86f0196..f9faaa93 100644
--- a/test/python/mocks.py
+++ b/test/python/mocks.py
@@ -49,6 +49,13 @@ class MockWordTable:
         self.conn.commit()
 
 
+    def add_country(self, country_code, word_token):
+        with self.conn.cursor() as cur:
+            cur.execute("INSERT INTO word (word_token, country_code) VALUES(%s, %s)",
+                        (word_token, country_code))
+        self.conn.commit()
+
+
     def add_postcode(self, word_token, postcode):
         with self.conn.cursor() as cur:
             cur.execute("""INSERT INTO word (word_token, word, class, type)
@@ -71,7 +78,18 @@ class MockWordTable:
         with self.conn.cursor() as cur:
             cur.execute("""SELECT word_token, word, class, type, operator
                            FROM word WHERE class != 'place'""")
-            return set((tuple(row) for row in cur))
+            result = set((tuple(row) for row in cur))
+            assert len(result) == cur.rowcount, "Word table has duplicates."
+            return result
+
+
+    def get_country(self):
+        with self.conn.cursor() as cur:
+            cur.execute("""SELECT country_code, word_token
+                           FROM word WHERE country_code is not null""")
+            result = set((tuple(row) for row in cur))
+            assert len(result) == cur.rowcount, "Word table has duplicates."
+            return result
 
 
     def get_postcodes(self):
@@ -80,6 +98,13 @@ class MockWordTable:
                            WHERE class = 'place' and type = 'postcode'""")
             return set((row[0] for row in cur))
 
+    def get_partial_words(self):
+        with self.conn.cursor() as cur:
+            cur.execute("""SELECT word_token, search_name_count FROM word
+                           WHERE class is null and country_code is null
+                                 and not word_token like ' %'""")
+            return set((tuple(row) for row in cur))
+
 
 class MockPlacexTable:
     """ A placex table for testing.