Created better environment for testing

scripts/create_database.py

@@ -16,9 +16,15 @@ def create_connection(path: str, filename: str):
         # Execute creation scripts
         sql = open("scripts/create_tables.sql", "r");
         conn.executescript(sql.read())
-        
-        print("SQLite Version: " + sqlite3.version)
+
         print("Table creation complete")
+
+        # Populate with test data if we are in Test Mode
+        if os.environ.get("ENVIRON") != "test":
+            return
+        
+        sql = open("scripts/test_data.sql", "r");
+        conn.executescript(sql.read())
     except sqlite3.Error as e:
         print(e)
     finally:
@@ -32,5 +38,20 @@ def create_directory(dir: str):
     except FileExistsError:
         pass
 
+def remove_file(dir: str):
+    try:
+        os.remove(dir)
+    except FileNotFoundError:
+        pass
 if __name__ == '__main__':
-    create_connection(r"./data/",  r"wmgzon.db")
+    dir = r"./data/"
+    db_name = r"wmgzon.db"
+
+    # Check for test environ
+    if os.environ.get("ENVIRON") == "test":
+        # Remove the original test database
+        print("TEST ENVIRONMENT ACTIVE")
+        db_name = "test_" + db_name
+        remove_file(dir + db_name)
+
+    create_connection(dir, db_name)