Repository: robotpt/interaction-engine

Issue:
Replace pickle database
Replace `pickle-database` with a JSON file or another database service. The current setup is bad for the following reasons:
* Not human readable
* Requires setting variables in a Python script

Discussion:

Patch:
diff --git a/example.py b/example.py
index 0f2a802..870bd94 100644
--- a/example.py
+++ b/example.py
@@ -1,8 +1,7 @@
 import os
 
-from pickled_database import PickledDatabase
-
 from interaction_engine import InteractionEngine
+from interaction_engine.json_database import Database
 from interaction_engine.messager import Message, Node, DirectedGraph
 from interaction_engine.planner import MessagerPlanner
 from interaction_engine.text_populator import TextPopulator, DatabasePopulator, VarietyPopulator
@@ -16,11 +15,11 @@ CREATE RESOURCES USED TO POPULATE TEXT
 import atexit
 
 # Create the database
-db_file = 'test_db.pkl'
-db = PickledDatabase(db_file)
-db.create_key_if_not_exists('user_name')
-db.create_key_if_not_exists('question_idx', 0)
-db.create_key_if_not_exists('answers')
+db_file = 'test_db.json'
+db = Database(db_file)
+db["user_name"] = None
+db["question_idx"] = 0
+db["answers"] = None
 
 # Create a file with text used for variation
 variation_file = 'variation.csv'
@@ -156,7 +155,7 @@ closing = DirectedGraph(
 ORGANIZE CONTENT TO BE PLAYED
 """
 
-interface_ = TerminalClientAndServerInterface(pickled_database=db)
+interface_ = TerminalClientAndServerInterface(database=db)
 graphs_ = [greeting, basic_questions, psych_question, closing]
 
 # Create a plan
diff --git a/interaction_engine/interfaces/client_and_server_interface.py b/interaction_engine/interfaces/client_and_server_interface.py
index ffcb809..8c05996 100644
--- a/interaction_engine/interfaces/client_and_server_interface.py
+++ b/interaction_engine/interfaces/client_and_server_interface.py
@@ -1,6 +1,6 @@
 from interaction_engine.messager.message import Message
 from interaction_engine.interfaces.interface import Interface
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 
 
 class ClientAndServerInterface(Interface):
@@ -11,7 +11,7 @@ class ClientAndServerInterface(Interface):
             multiple_choice_input_fn,
             direct_input_output_fn,
             direct_input_input_fn,
-            pickled_database=None,
+            database=None,
             is_create_db_key_if_not_exist=False,
     ):
         super().__init__(
@@ -27,16 +27,16 @@ class ClientAndServerInterface(Interface):
                     Message.Type.DIRECT_INPUT: direct_input_output_fn,
                 }
             ),
-            pickled_database=pickled_database,
+            database=database,
             is_create_db_key_if_not_exist=is_create_db_key_if_not_exist
         )
 
         if (
-                pickled_database is not None
-                and not issubclass(pickled_database.__class__, PickledDatabase)
+                database is not None
+                and not type(database) == Database
         ):
             raise TypeError
-        self._db = pickled_database
+        self._db = database
         self._is_create_db_key_if_not_exist = is_create_db_key_if_not_exist
 
     @staticmethod
diff --git a/interaction_engine/interfaces/interface.py b/interaction_engine/interfaces/interface.py
index 7b20459..cbe8e45 100644
--- a/interaction_engine/interfaces/interface.py
+++ b/interaction_engine/interfaces/interface.py
@@ -1,6 +1,6 @@
 from interaction_engine.messager.message import Message
 from robotpt_common_utils import lists, user_input
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 
 
 class Interface:
@@ -9,7 +9,7 @@ class Interface:
             self,
             input_fn,
             output_fn=None,
-            pickled_database=None,
+            database=None,
             is_create_db_key_if_not_exist=False,
     ):
         self._input_fn = input_fn
@@ -18,11 +18,11 @@ class Interface:
         self._output_fn = output_fn
 
         if (
-                pickled_database is not None
-                and not issubclass(pickled_database.__class__, PickledDatabase)
+                database is not None
+                and not type(database) is Database
         ):
             raise TypeError
-        self._db = pickled_database
+        self._db = database
         self._is_create_db_key_if_not_exist = is_create_db_key_if_not_exist
 
     def run(self, message):
@@ -64,18 +64,19 @@ class Interface:
         ):
             if key not in self._db:
                 if self._is_create_db_key_if_not_exist:
-                    self._db.create_key(key)
+                    self._db[key] = None
                 else:
                     raise KeyError(f"'{key}' doesn't exist in the database")
 
             if message.is_append_result:
                 save_result = [result]
                 if self._db.is_set(key):
-                    save_result = self._db.get(key) + save_result
+
+                    save_result = self._db[key] + save_result
             else:
                 save_result = result
 
-            self._db.set(message.result_db_key, save_result)
+            self._db[message.result_db_key] = save_result
 
         return result
 
diff --git a/interaction_engine/interfaces/terminal_client_and_server_interface.py b/interaction_engine/interfaces/terminal_client_and_server_interface.py
index 504d9b9..e4c3336 100644
--- a/interaction_engine/interfaces/terminal_client_and_server_interface.py
+++ b/interaction_engine/interfaces/terminal_client_and_server_interface.py
@@ -1,7 +1,7 @@
 from interaction_engine.interfaces.client_and_server_interface import ClientAndServerInterface
 from interaction_engine.messager.message import Message
+from interaction_engine.json_database import Database
 from robotpt_common_utils import math_tools, lists
-from pickled_database import PickledDatabase
 import textwrap
 
 
@@ -9,7 +9,7 @@ class TerminalClientAndServerInterface(ClientAndServerInterface):
 
     def __init__(
             self,
-            pickled_database=None,
+            database=None,
             is_create_db_key_if_not_exist=True,
     ):
         fns = [
@@ -20,7 +20,7 @@ class TerminalClientAndServerInterface(ClientAndServerInterface):
         ]
         super().__init__(
             *fns,
-            pickled_database=pickled_database,
+            database=database,
             is_create_db_key_if_not_exist=is_create_db_key_if_not_exist
         )
 
@@ -111,11 +111,11 @@ if __name__ == '__main__':
     import os
     import atexit
 
-    db_file = "memory.pkl"
+    db_file = "memory.json"
     atexit.register(lambda: os.remove(db_file))
 
-    db = PickledDatabase(db_file)
-    interface = TerminalClientAndServerInterface(pickled_database=db)
+    db = Database(db_file)
+    interface = TerminalClientAndServerInterface(database=db)
     for msg in [
         multiple_choice_message1,
         multiple_choice_message2,
diff --git a/interaction_engine/json_database/__init__.py b/interaction_engine/json_database/__init__.py
new file mode 100644
index 0000000..630a6b0
--- /dev/null
+++ b/interaction_engine/json_database/__init__.py
@@ -0,0 +1,1 @@
+from interaction_engine.json_database.database import Database
diff --git a/interaction_engine/json_database/database.py b/interaction_engine/json_database/database.py
new file mode 100644
index 0000000..880d189
--- /dev/null
+++ b/interaction_engine/json_database/database.py
@@ -0,0 +1,142 @@
+import json
+import logging
+import os
+import pprint
+
+logging.basicConfig(level=logging.INFO)
+
+
+class Database:
+
+    def __init__(
+            self,
+            database_file_name
+    ):
+
+        self._database_file = self._make_sure_database_file_is_valid(
+            database_file_name
+        )
+
+        if not os.path.exists(database_file_name):
+            self._database = {}
+            self._create_new_database_file(self._database_file)
+
+        self._database = self._load_database_from_file()
+
+    def _create_new_database_file(self, database_file_name):
+        directory = os.path.dirname(database_file_name)
+        if directory != "":
+            os.makedirs(directory, exist_ok=True)
+        with open(self._database_file, "w"):
+            self.save_to_database_file()
+
+    def _load_database_from_file(self):
+        with open(self._database_file) as f:
+            try:
+                database = json.load(f)
+            except json.JSONDecodeError:
+                logging.info("Error decoding JSON file, defaulting to empty database")
+                database = {}
+        return database
+
+    def _make_sure_database_file_is_valid(self, database_file_name):
+        if type(database_file_name) is not str:
+            raise TypeError("Not a valid file name.")
+
+        if not database_file_name.endswith(".json"):
+            raise TypeError("Must be a .json file.")
+
+        return database_file_name
+
+    def __contains__(self, item):
+        self._database = self._load_database_from_file()
+        return item in self._database
+
+    def __repr__(self):
+        return pprint.pformat(self._database)
+
+    def __getitem__(self, item):
+        self._database = self._load_database_from_file()
+        return self._database[item]
+
+    def __setitem__(self, key, value):
+        self._database = self._load_database_from_file()
+        possible_value_types = [list, str, int, float]
+        if type(key) is not str:
+            raise TypeError("Key must be a string.")
+
+        if value is None:
+            value = ""
+        elif type(value) in possible_value_types:
+            if type(value) is list:
+                self._check_if_items_in_list_are_valid(value, [str, int, float])
+        else:
+            raise TypeError("Invalid value: {value}, can't have type {type}".format(value=value, type=type(value)))
+
+        self._database[key] = value
+        self.save_to_database_file()
+
+    def _check_if_items_in_list_are_valid(self, list, possible_value_types):
+        for item in list:
+            if type(item) not in possible_value_types:
+                raise TypeError("List contains invalid value: {item}, type {type}".format(
+                    item=item,
+                    type=type(item)
+                ))
+
+    def clear_value(self, key):
+        self._database = self._load_database_from_file()
+        self._database[key] = ""
+        self.save_to_database_file()
+
+    def clear_entire_database(self):
+        for key in self._database:
+            self.clear_value(key)
+
+    def create_key(self, key):
+        self._database = self._load_database_from_file()
+        self._database[key] = ""
+        self.save_to_database_file()
+
+    def delete_database_file(self):
+        os.remove(self._database_file)
+
+    def delete_key(self, key):
+        try:
+            self._database.pop(key)
+        except KeyError as e:
+            logging.info(e)
+            pass
+        self.save_to_database_file()
+
+    def get_keys(self):
+        self._database = self._load_database_from_file()
+        return [i for i in self._database.keys()]
+
+    def is_set(self, key):
+        self._database = self._load_database_from_file()
+        return self._database[key] != ""
+
+    def save_to_database_file(self):
+        with open(self._database_file, "w") as f:
+            json.dump(self._database, f)
+
+    def reset_database(self):
+        self._database = {}
+        self.save_to_database_file()
+
+    @property
+    def database(self):
+        self._database = self._load_database_from_file()
+        return self._database
+
+    @database.setter
+    def database(self, database):
+        if type(database) is not dict:
+            raise TypeError("Database must be a dict.")
+        self._database = database
+        self.save_to_database_file()
+
+    @property
+    def database_file(self):
+        return self._database_file
diff --git a/interaction_engine/messager/directed_graph.py b/interaction_engine/messager/directed_graph.py
index 5c3fc07..a9dc4ea 100644
--- a/interaction_engine/messager/directed_graph.py
+++ b/interaction_engine/messager/directed_graph.py
@@ -94,18 +94,18 @@ class DirectedGraph(BaseMessenger):
 if __name__ == '__main__':
 
     import os
-    from pickled_database import PickledDatabase
+    from interaction_engine.json_database import Database
     from interaction_engine.text_populator import TextPopulator
     from interaction_engine.text_populator import DatabasePopulator
     from interaction_engine.text_populator.variety_populator import VarietyPopulator
 
-    db_file = 'test_db.pkl'
-    db = PickledDatabase(db_file)
-    db.create_key_if_not_exists('key1', 1)
-    db.create_key_if_not_exists('key2', 'two')
-    db.create_key_if_not_exists('no_value_key')
-    db.create_key_if_not_exists('user_name', 'Audrow')
-    db.create_key_if_not_exists('question_idx', 1)
+    db_file = 'test_db.json'
+    db = Database(db_file)
+    db['key1'] = 1
+    db['key2'] = 'two'
+    db['no_value_key'] = None
+    db['user_name'] = 'Audrow'
+    db['question_idx'] = 1
 
     variation_file = 'variation.csv'
     variation_file_contents = """
@@ -187,9 +187,8 @@ fakebar,fake-bar
     )
 
     from interaction_engine.interfaces.terminal_client_and_server_interface import TerminalClientAndServerInterface
-    from pickled_database import PickledDatabase
 
-    interface = TerminalClientAndServerInterface(pickled_database=db)
+    interface = TerminalClientAndServerInterface(database=db)
     while directed_graph.is_active:
         msg = directed_graph.get_message()
         user_response = interface.run(msg)
diff --git a/interaction_engine/messager/message.py b/interaction_engine/messager/message.py
index f6a17af..d24bf92 100644
--- a/interaction_engine/messager/message.py
+++ b/interaction_engine/messager/message.py
@@ -185,17 +185,17 @@ class Message(BaseMessenger):
 if __name__ == "__main__":
 
     import os
-    from pickled_database import PickledDatabase
+    from interaction_engine.json_database import Database
     from interaction_engine.text_populator import DatabasePopulator
     from interaction_engine.text_populator.variety_populator import VarietyPopulator
 
-    db_file = 'test_db.pkl'
-    db = PickledDatabase(db_file)
-    db.create_key_if_not_exists('key1', 1)
-    db.create_key_if_not_exists('key2', 'two')
-    db.create_key_if_not_exists('no_value_key')
-    db.create_key_if_not_exists('user_name', 'Audrow')
-    db.create_key_if_not_exists('question_idx', 1)
+    db_file = 'test_db.json'
+    db = Database(db_file)
+    db['key1'] = 1
+    db['key2'] = 'two'
+    db['no_value_key'] = None
+    db['user_name'] = 'Audrow'
+    db['question_idx'] = 1
 
     my_str = """
 {'var': 'greeting'}, {'db': 'user_name'}. 
diff --git a/interaction_engine/text_populator/database_populator.py b/interaction_engine/text_populator/database_populator.py
index f47776b..9a546e1 100644
--- a/interaction_engine/text_populator/database_populator.py
+++ b/interaction_engine/text_populator/database_populator.py
@@ -1,8 +1,9 @@
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 from interaction_engine.text_populator.base_populator import BasePopulator
 
 import datetime
 
+
 class DatabasePopulator(BasePopulator):
 
     class Tags:
@@ -25,7 +26,7 @@ class DatabasePopulator(BasePopulator):
         )
 
         if type(database) is str:
-            self._db = PickledDatabase(database)
+            self._db = Database(database)
         else:
             self._db = database
 
@@ -41,7 +42,7 @@ class DatabasePopulator(BasePopulator):
             modify_before_resaving_fn=None,
     ):
         if self._db.is_set(key):
-            value = self._db.get(key)
+            value = self._db[key]
         elif default_value is not None:
             value = default_value
         else:
@@ -51,7 +52,7 @@ class DatabasePopulator(BasePopulator):
             if modify_before_resaving_fn in DatabasePopulator._common_fns:
                 modify_before_resaving_fn = DatabasePopulator._common_fns[modify_before_resaving_fn]
             new_value = modify_before_resaving_fn(value)
-            self._db.set(key, new_value)
+            self._db[key] = new_value
 
         if type(value) is datetime.time:
             value = value.strftime("%-I:%M%p")
diff --git a/setup.py b/setup.py
index fb73c0c..1768e05 100644
--- a/setup.py
+++ b/setup.py
@@ -22,7 +22,6 @@ VERSION = '0.0.6'
 
 # What packages are required for this module to be executed?
 REQUIRED = [
-    'pickle_database',
     'robotpt_common_utils',
     # 'requests', 'maya', 'records',
 ]


Test patch:
diff --git a/tests/test_database.py b/tests/test_database.py
new file mode 100644
index 0000000..a1341c6
--- /dev/null
+++ b/tests/test_database.py
@@ -0,0 +1,98 @@
+import atexit
+import json
+import mock
+import os
+import unittest
+
+from interaction_engine.json_database import Database
+
+
+class TestDatabase(unittest.TestCase):
+
+    def setUp(self):
+        self._database_file_name = "test_database.json"
+        test_database = {
+            "key1": "value1",
+            "key2": "value2"
+        }
+        with open(self._database_file_name, "w") as f:
+            json.dump(test_database, f)
+
+    def tearDown(self) -> None:
+        os.remove(self._database_file_name)
+
+    def test_create_new_database_file(self):
+        database = Database(self._database_file_name)
+
+        self.assertEqual(
+            self._database_file_name,
+            database.database_file,
+            os.path.join(
+                os.path.dirname(os.path.realpath(__file__)),
+            )
+        )
+
+    def test_load_from_database_file(self):
+        expected_database = {
+            "key1": "value1",
+            "key2": "value2"
+        }
+        with open(self._database_file_name, "w") as f:
+            json.dump(expected_database, f)
+        database = Database(self._database_file_name)
+        self.assertEqual(expected_database, database.database)
+
+    def test_load_from_empty_database_file(self):
+        expected_database = {}
+        with open(self._database_file_name, "w") as f:
+            json.dump(expected_database, f)
+        database = Database("test_database.json")
+        self.assertEqual({}, database.database)
+
+    def test_set_database(self):
+        database = Database(self._database_file_name)
+        database["key3"] = "value3"
+        self.assertEqual("value3", database["key3"])
+
+        database["key4"] = "value4"
+        self.assertEqual("value4", database["key4"])
+
+    def test_clear_database(self):
+        expected_initial_database = {"key1": "value1", "key2": "value2"}
+        expected_cleared_database = {"key1": "", "key2": ""}
+
+        database = Database(self._database_file_name)
+        self.assertEqual(expected_initial_database, database.database)
+
+        database.clear_value("key1")
+        database.clear_value("key2")
+        self.assertEqual(expected_cleared_database, database.database)
+
+        database.database = expected_initial_database
+        self.assertEqual(expected_initial_database, database.database)
+
+        database.clear_entire_database()
+        self.assertEqual(expected_cleared_database, database.database)
+
+    def test_access_same_file_from_multiple_databases(self):
+        test_database = {"key1": "value1", "key2": "value2"}
+
+        database1 = Database(self._database_file_name)
+        database2 = Database(self._database_file_name)
+
+        self.assertEqual(test_database, database1.database)
+        self.assertEqual(test_database, database2.database)
+
+        database1["key3"] = "value3"
+
+        expected_database = {
+            "key1": "value1",
+            "key2": "value2",
+            "key3": "value3",
+        }
+        self.assertEqual(expected_database, database1.database)
+        self.assertEqual(expected_database, database2.database)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_database_populator.py b/tests/test_database_populator.py
index ce4efac..0761b01 100644
--- a/tests/test_database_populator.py
+++ b/tests/test_database_populator.py
@@ -2,17 +2,18 @@ import unittest
 import os
 
 from interaction_engine.text_populator.database_populator import DatabasePopulator
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 
 
 class TestDatabasePopulator(unittest.TestCase):
 
     def setUp(self):
-        self.file = 'test_db.pkl'
-        db = PickledDatabase(self.file)
-        db.create_key('key1', 1)
-        db.create_key('key2', 'two')
-        db.create_key('no_value_key')
+        self.file = 'test_db.json'
+        db = Database(self.file)
+        db['key1'] = '1'
+        db['key2'] = 'two'
+        db['no_value_key'] = None
+        db['key3'] = 0
         self.dp = DatabasePopulator(db)
 
     def tearDown(self):
@@ -47,3 +48,17 @@ class TestDatabasePopulator(unittest.TestCase):
             'not_a_key'
         )
 
+    def test_post_op(self):
+        for i in range(10):
+            self.assertEqual(
+                i,
+                self.dp._db['key3']
+            )
+            self.dp.get_replacement('key3', modify_before_resaving_fn='increment')
+
+        for i in range(10):
+            self.assertEqual(
+                10-i,
+                self.dp._db['key3']
+            )
+            self.dp.get_replacement('key3', modify_before_resaving_fn='decrement')
diff --git a/tests/test_interface.py b/tests/test_interface.py
index 6cb2723..394f646 100644
--- a/tests/test_interface.py
+++ b/tests/test_interface.py
@@ -1,28 +1,28 @@
 from interaction_engine.interfaces.client_and_server_interface import ClientAndServerInterface
 from interaction_engine.messager import Message
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 from robotpt_common_utils import lists
 
 import unittest
 import os
 
-db_file = 'test_db.pkl'
+db_file = 'test_db.json'
 
 
 class TestInterface(unittest.TestCase):
 
     def setUp(self) -> None:
 
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         self._db = db
         self._io = TestIO()
         self._interface = ClientAndServerInterface(
             *[self._io.output_fn, self._io.input_fn]*2,
-            pickled_database=self._db
+            database=self._db
         )
 
     def tearDown(self) -> None:
@@ -169,7 +169,7 @@ class TestInterface(unittest.TestCase):
         )
         self.assertEqual(
             choice,
-            self._db.get(db_key)
+            self._db[db_key]
         )
 
     def test_append_db_entry(self):
@@ -205,7 +205,7 @@ class TestInterface(unittest.TestCase):
         )
 
         truth_out = choices
-        out = self._db.get(db_key)
+        out = self._db.create_key(db_key)
         for i in range(len(truth_out)):
             self.assertEqual(
                 truth_out[i],
@@ -238,7 +238,7 @@ class TestInterface(unittest.TestCase):
         )
         self.assertEqual(
             choice,
-            self._db.get(db_key)
+            self._db[db_key]
         )
 
 class TestIO:
diff --git a/tests/test_message.py b/tests/test_message.py
index 7d2c2b6..ad80dc3 100644
--- a/tests/test_message.py
+++ b/tests/test_message.py
@@ -4,9 +4,9 @@ import random
 
 from interaction_engine.messager import Message
 from interaction_engine.text_populator import TextPopulator, VarietyPopulator, DatabasePopulator
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 
-db_file = 'test_db.pkl'
+db_file = 'test_db.json'
 
 variation_file = 'variation.csv'
 variation_file_contents = """
@@ -32,10 +32,10 @@ class TestMessage(unittest.TestCase):
 
     def setUp(self) -> None:
 
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         with open(variation_file, 'w', newline='') as csvfile:
             csvfile.write(variation_file_contents.strip())
diff --git a/tests/test_populator.py b/tests/test_populator.py
index 1785860..49a5f30 100644
--- a/tests/test_populator.py
+++ b/tests/test_populator.py
@@ -1,11 +1,12 @@
+import json
 import unittest
 import os
-from pickled_database import PickledDatabase
+from interaction_engine.json_database import Database
 from interaction_engine.text_populator import TextPopulator
 from interaction_engine.text_populator.database_populator import DatabasePopulator
 from interaction_engine.text_populator.variety_populator import VarietyPopulator
 
-db_file = 'test_db.pkl'
+db_file = 'test_db.json'
 
 variation_file = 'variation.csv'
 variation_file_contents = """
@@ -32,10 +33,10 @@ class TestTextPopulator(unittest.TestCase):
 
     def setUp(self) -> None:
 
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         with open(variation_file, 'w', newline='') as csvfile:
             csvfile.write(variation_file_contents.strip())
@@ -218,12 +219,12 @@ class TestTextPopulator(unittest.TestCase):
                 self._text_populator.is_valid(test_str)
             self.assertEqual(
                 i,
-                self._db.get('question_idx')
+                self._db["question_idx"]
             )
             self._text_populator.run(test_str)
             self.assertEqual(
                 i+1,
-                self._db.get('question_idx')
+                self._db["question_idx"]
             )
 
         test_str = "{'var': '{'var': 'foo'}bar'}"
@@ -258,10 +259,10 @@ class TestTextPopulator(unittest.TestCase):
         greeting,Hi
         greeting,Hello{'var': 'foo'}
         """
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         with open(bad_variation_file, 'w', newline='') as csvfile:
             csvfile.write(bad_variation_file_contents.strip())
@@ -284,10 +285,10 @@ class TestTextPopulator(unittest.TestCase):
         greeting,Hi
         greeting,Hello{
         """
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         with open(bad_variation_file, 'w', newline='') as csvfile:
             csvfile.write(bad_variation_file_contents.strip())
@@ -310,10 +311,10 @@ class TestTextPopulator(unittest.TestCase):
         greeting,Hi ready for number {'db': 'question_idx', 'post-op': 'increment'}
         greeting,Hello {'db': 'user_name'}
         """
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         with open(bad_variation_file, 'w', newline='') as csvfile:
             csvfile.write(bad_variation_file_contents.strip())
@@ -334,10 +335,10 @@ class TestTextPopulator(unittest.TestCase):
         greeting,"Hi ready for number {'db': 'question_idx', 'post-op': 'increment'}"
         greeting,Hello {'db': 'user_name'}
         """
-        db = PickledDatabase(db_file)
-        db.create_key_if_not_exists('user_name')
-        db.create_key_if_not_exists('question_idx', 0)
-        db.create_key_if_not_exists('answers')
+        db = Database(db_file)
+        db["user_name"] = None
+        db["question_idx"] = 0
+        db["answers"] = None
 
         with open(correct_variation_file, 'w', newline='') as csvfile:
             csvfile.write(correct_variation_file_contents.strip())
@@ -358,7 +359,7 @@ class TestTextPopulator(unittest.TestCase):
         greeting,Hi
         greeting,Hello {'db': 'not_a_key'}
         """
-        db = PickledDatabase(db_file)
+        db = Database(db_file)
 
         with open(bad_variation_file, 'w', newline='') as csvfile:
             csvfile.write(bad_variation_file_contents.strip())