X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/b1d33e6b4948b37e4686d80dc2336c08dc0e7548..bc981d0261896e240f8cd4cf212f7d191a71e152:/test/python/mocks.py

diff --git a/test/python/mocks.py b/test/python/mocks.py
index d86f0196..e95d5772 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):