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_full_word(self, word_id, word, word_token=None):
  21         with self.conn.cursor() as cur:
  22             cur.execute("""INSERT INTO word (word_id, word_token, type, word, info)
  23                            VALUES(%s, %s, 'W', %s, '{}'::jsonb)""",
  24                         (word_id, word or word_token, word))
  25         self.conn.commit()
  26
  27
  28     def add_special(self, word_token, word, cls, typ, oper):
  29         with self.conn.cursor() as cur:
  30             cur.execute("""INSERT INTO word (word_token, type, word, info)
  31                               VALUES (%s, 'S', %s,
  32                                       json_build_object('class', %s,
  33                                                         'type', %s,
  34                                                         'op', %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, type, word)
  42                            VALUES(%s, 'C', %s)""",
  43                         (word_token, country_code))
  44         self.conn.commit()
  45
  46
  47     def add_postcode(self, word_token, postcode):
  48         with self.conn.cursor() as cur:
  49             cur.execute("""INSERT INTO word (word_token, type, word)
  50                               VALUES (%s, 'P', %s)
  51                         """, (word_token, postcode))
  52         self.conn.commit()
  53
  54
  55     def count(self):
  56         with self.conn.cursor() as cur:
  57             return cur.scalar("SELECT count(*) FROM word")
  58
  59
  60     def count_special(self):
  61         with self.conn.cursor() as cur:
  62             return cur.scalar("SELECT count(*) FROM word WHERE type = 'S'")
  63
  64
  65     def get_special(self):
  66         with self.conn.cursor() as cur:
  67             cur.execute("SELECT word_token, info, word FROM word WHERE type = 'S'")
  68             result = set(((row[0], row[2], row[1]['class'],
  69                            row[1]['type'], row[1]['op']) for row in cur))
  70             assert len(result) == cur.rowcount, "Word table has duplicates."
  71             return result
  72
  73
  74     def get_country(self):
  75         with self.conn.cursor() as cur:
  76             cur.execute("SELECT word, word_token FROM word WHERE type = 'C'")
  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 WHERE type = 'P'")
  85             return set((row[0] for row in cur))
  86
  87
  88     def get_partial_words(self):
  89         with self.conn.cursor() as cur:
  90             cur.execute("SELECT word_token, info FROM word WHERE type ='w'")
  91             return set(((row[0], row[1]['count']) for row in cur))
  92