Database insert tests (#153)

fixes #48 
* mongodb: test if two same-key insertions update body

* change mongodb insert to upsert

* skip mongodb test when running via travis

* sqlitedb: test if two same-key insertions update body

* sqlite: add upsert

pastepwn/database/mongodb.py

@@ -28,7 +28,7 @@ def __init__(self, ip="127.0.0.1", port=27017, dbname="pastepwn", collectionname
         self.collection.create_index([('key', pymongo.ASCENDING)], unique=True)
 
     def _insert_data(self, data):
-        self.collection.insert_one(data)
+        self.collection.update_one({'key': data['key']}, {'$set': data}, upsert=True)
 
     def _get_data(self, key, value):
         return self.collection.find({key: value})

pastepwn/database/sqlitedb.py

@@ -63,6 +63,12 @@ def _insert_data(self, paste):
                              paste.body))
         self.db.commit()
 
+    def _update_data(self, paste):
+        self.cursor.execute("UPDATE pastes SET body = ? WHERE key = ?",
+                            (paste.body,
+                            paste.key))
+        self.db.commit()
+
     def _get_data(self, key, value):
         pass
 
@@ -80,6 +86,9 @@ def store(self, paste):
             self._insert_data(paste)
         except Exception as e:
             self.logger.debug("Exception '{0}'".format(e))
+            if "UNIQUE constraint failed: pastes.key" in str(e):
+                self.logger.debug("Doing upsert")
+                self._update_data(paste)
 
     def get(self, key):
         return self._get_data("key", key)

pastepwn/database/tests/mongodb_test.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+
+import os
+import random
+import string
+import unittest
+
+from pastepwn import Paste
+from pastepwn.database import MongoDB
+
+
+@unittest.skipIf("TRAVIS" in os.environ and os.environ["TRAVIS"] == "true", "Skipping this test on Travis CI.")
+class MongoDBTest(unittest.TestCase):
+
+    def setUp(self):
+        rand_text = list()
+        for i in range(3):
+            rand_text.append(''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(8)))
+
+        p = {"scrape_url": "https://scrape.pastebin.com/api_scrape_item.php?i=" + rand_text[0],
+             "full_url": "https://pastebin.com/" + rand_text[0],
+             "date": "1442911802",
+             "key": rand_text[0],
+             "size": "890",
+             "expire": "1442998159",
+             "title": "Once we all know when we goto function",
+             "syntax": "java",
+             "user": "admin",
+             "body": rand_text[1:],
+             }
+
+        self.p = p
+        self.paste = Paste(p.get("key"),
+                           p.get("title"),
+                           p.get("user"),
+                           p.get("size"),
+                           p.get("date"),
+                           p.get("expire"),
+                           p.get("syntax"),
+                           p.get("scrape_url"),
+                           p.get("full_url"))
+
+        self.database = MongoDB(collectionname="pastepwn_test")
+
+    def tearDown(self):
+        self.database.db.drop_collection("pastepwn_test")
+
+    def test_insert_same_key(self):
+        for body_text in self.p['body']:
+            self.paste.set_body(body_text)
+            self.database.store(self.paste)
+
+        self.assertEqual(self.database.get(self.p['key']).next()['body'], self.p['body'][1])

pastepwn/database/tests/sqlite_test.py

@@ -0,0 +1,53 @@
+# -*- coding: utf-8 -*-
+
+import os
+import random
+import shutil
+import string
+import unittest
+
+from pastepwn import Paste
+from pastepwn.database import SQLiteDB
+
+
+class SQLiteDBTest(unittest.TestCase):
+
+    def setUp(self):
+        rand_text = list()
+        for i in range(3):
+            rand_text.append(''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for _ in range(8)))
+
+        p = {"scrape_url": "https://scrape.pastebin.com/api_scrape_item.php?i=" + rand_text[0],
+             "full_url": "https://pastebin.com/" + rand_text[0],
+             "date": "1442911802",
+             "key": rand_text[0],
+             "size": "890",
+             "expire": "1442998159",
+             "title": "Once we all know when we goto function",
+             "syntax": "java",
+             "user": "admin",
+             "body": rand_text[1:],
+             }
+
+        self.p = p
+        self.paste = Paste(p.get("key"),
+                           p.get("title"),
+                           p.get("user"),
+                           p.get("size"),
+                           p.get("date"),
+                           p.get("expire"),
+                           p.get("syntax"),
+                           p.get("scrape_url"),
+                           p.get("full_url"))
+
+        self.database = SQLiteDB(dbpath="sqlite_test/pastepwn_test")
+
+    def tearDown(self):
+        shutil.rmtree(os.path.dirname(self.database.dbpath))
+
+    def test_insert_same_key(self):
+        for body_text in self.p['body']:
+            self.paste.set_body(body_text)
+            self.database.store(self.paste)
+
+        self.assertEqual(self.database.cursor.execute("SELECT body FROM pastes WHERE key = \'{0}\'".format(self.p['key'])).fetchone()[0], self.p['body'][1])