Merge pull request #2360 from AntoJvlt/postcodes-place-table

[nominatim.git] / test / python / mocks.py

diff --git a/test/python/mocks.py b/test/python/mocks.py
index 53d84aace63b2783ac87c63f2be2bda389324696..e95d57721357b7c5c8e7896feaf6434248f98b6f 100644 (file)
--- a/test/python/mocks.py
+++ b/test/python/mocks.py
@@ -5,6 +5,8 @@ import itertools
 
 import psycopg2.extras
 
+from nominatim.db import properties
+
 class MockParamCapture:
     """ Mock that records the parameters with which a function was called
         as well as the number of calls.
@@ -12,6 +14,8 @@ class MockParamCapture:
     def __init__(self, retval=0):
         self.called = 0
         self.return_value = retval
+        self.last_args = None
+        self.last_kwargs = None
 
     def __call__(self, *args, **kwargs):
         self.called += 1
@@ -37,11 +41,18 @@ class MockWordTable:
 
         conn.commit()
 
-    def add_special(self, word_token, word, cls, typ, op):
+    def add_special(self, word_token, word, cls, typ, oper):
         with self.conn.cursor() as cur:
             cur.execute("""INSERT INTO word (word_token, word, class, type, operator)
                               VALUES (%s, %s, %s, %s, %s)
-                        """, (word_token, word, cls, typ, op))
+                        """, (word_token, word, cls, typ, oper))
+        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()
 
 
@@ -67,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):
@@ -125,3 +147,16 @@ class MockPlacexTable:
                          admin_level, address, extratags, 'SRID=4326;' + geom,
                          country))
         self.conn.commit()
+
+
+class MockPropertyTable:
+    """ A property table for testing.
+    """
+    def __init__(self, conn):
+        self.conn = conn
+
+
+    def set(self, name, value):
+        """ Set a property in the table to the given value.
+        """
+        properties.set_property(self.conn, name, value)