test/python/mock_legacy_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 MockLegacyWordTable:
   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,
  14                                               word text,
  15                                               class text,
  16                                               type text,
  17                                               country_code varchar(2),
  18                                               search_name_count INTEGER,
  19                                               operator TEXT)""")
  20
  21         conn.commit()
  22
  23     def add_full_word(self, word_id, word, word_token=None):
  24         with self.conn.cursor() as cur:
  25             cur.execute("""INSERT INTO word (word_id, word_token, word)
  26                            VALUES (%s, %s, %s)
  27                         """, (word_id, ' ' + (word_token or word), word))
  28         self.conn.commit()
  29
  30
  31     def add_special(self, word_token, word, cls, typ, oper):
  32         with self.conn.cursor() as cur:
  33             cur.execute("""INSERT INTO word (word_token, word, class, type, operator)
  34                               VALUES (%s, %s, %s, %s, %s)
  35                         """, (word_token, word, cls, typ, oper))
  36         self.conn.commit()
  37
  38
  39     def add_country(self, country_code, word_token):
  40         with self.conn.cursor() as cur:
  41             cur.execute("INSERT INTO word (word_token, country_code) VALUES(%s, %s)",
  42                         (word_token, country_code))
  43         self.conn.commit()
  44
  45
  46     def add_postcode(self, word_token, postcode):
  47         with self.conn.cursor() as cur:
  48             cur.execute("""INSERT INTO word (word_token, word, class, type)
  49                               VALUES (%s, %s, 'place', 'postcode')
  50                         """, (word_token, postcode))
  51         self.conn.commit()
  52
  53
  54     def count(self):
  55         with self.conn.cursor() as cur:
  56             return cur.scalar("SELECT count(*) FROM word")
  57
  58
  59     def count_special(self):
  60         with self.conn.cursor() as cur:
  61             return cur.scalar("SELECT count(*) FROM word WHERE class != 'place'")
  62
  63
  64     def get_special(self):
  65         with self.conn.cursor() as cur:
  66             cur.execute("""SELECT word_token, word, class, type, operator
  67                            FROM word WHERE class != 'place'""")
  68             result = set((tuple(row) for row in cur))
  69             assert len(result) == cur.rowcount, "Word table has duplicates."
  70             return result
  71
  72
  73     def get_country(self):
  74         with self.conn.cursor() as cur:
  75             cur.execute("""SELECT country_code, word_token
  76                            FROM word WHERE country_code is not null""")
  77             result = set((tuple(row) for row in cur))
  78             assert len(result) == cur.rowcount, "Word table has duplicates."
  79             return result
  80
  81
  82     def get_postcodes(self):
  83         with self.conn.cursor() as cur:
  84             cur.execute("""SELECT word FROM word
  85                            WHERE class = 'place' and type = 'postcode'""")
  86             return set((row[0] for row in cur))
  87
  88     def get_partial_words(self):
  89         with self.conn.cursor() as cur:
  90             cur.execute("""SELECT word_token, search_name_count FROM word
  91                            WHERE class is null and country_code is null
  92                                  and not word_token like ' %'""")
  93             return set((tuple(row) for row in cur))
  94