From 982b9dcb472b917dbeb360b4ddbbda91d91c8cbe Mon Sep 17 00:00:00 2001
From: Bradlee Speice <bradlee.speice@gmail.com>
Date: Tue, 16 Aug 2016 17:21:09 -0400
Subject: [PATCH] Get some tests in place for the MongoDB stuff

---
 .travis.yml                       |  5 +++--
 conftest.py                       |  7 +++++++
 metrik/conf.py                    | 14 ++++++++++++-
 metrik/targets/mongo_target.py    |  1 -
 test/mongo_test.py                | 25 +++++++++++++++++++++++
 test/targets/test_mongo_target.py | 34 +++++++++++++++++++++++++++++++
 6 files changed, 82 insertions(+), 4 deletions(-)
 create mode 100644 conftest.py
 create mode 100644 test/mongo_test.py
 create mode 100644 test/targets/test_mongo_target.py

diff --git a/.travis.yml b/.travis.yml
index e04a26b..8c3f31e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -5,7 +5,8 @@ python:
   - "3.4"
   - "3.5"
 
-before_install:
-  - export TZ=America/New_York
+services:
+  - mongodb
+  
 install: pip install -r requirements.txt
 script: python setup.py test
\ No newline at end of file
diff --git a/conftest.py b/conftest.py
new file mode 100644
index 0000000..66520d4
--- /dev/null
+++ b/conftest.py
@@ -0,0 +1,7 @@
+def pytest_configure(config):
+    import sys
+    sys._called_from_test = True
+
+
+def pytest_unconfigure(config):
+    del sys._called_from_test
diff --git a/metrik/conf.py b/metrik/conf.py
index ae1d084..f874786 100644
--- a/metrik/conf.py
+++ b/metrik/conf.py
@@ -1,4 +1,16 @@
+import os
+import sys
+
+IS_TRAVIS = 'TRAVIS_BUILD_NUMBER' in os.environ
+IS_PYTEST = hasattr(sys, '_called_from_test')
+
+TEST = IS_PYTEST or IS_TRAVIS
+
 USER_AGENT = "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:47.0) Gecko/20100101 Firefox/47.0"
 MONGO_HOST = 'localhost'
 MONGO_PORT = 27017
-MONGO_DATABASE = 'metrik'
\ No newline at end of file
+
+if TEST:
+    MONGO_DATABASE = 'metrik'
+else:
+    MONGO_DATABASE = 'metrik-test'
diff --git a/metrik/targets/mongo_target.py b/metrik/targets/mongo_target.py
index 09c0005..8d73b3b 100644
--- a/metrik/targets/mongo_target.py
+++ b/metrik/targets/mongo_target.py
@@ -19,7 +19,6 @@ class MongoTarget(Target):
     def persist(self, dict_object):
         id_dict = dict_object
         id_dict['_id'] = self.id
-        id_dict['_retrieved'] = datetime.now()
         return self.collection.insert_one(id_dict).inserted_id
 
     def retrieve(self):
diff --git a/test/mongo_test.py b/test/mongo_test.py
new file mode 100644
index 0000000..c145a7a
--- /dev/null
+++ b/test/mongo_test.py
@@ -0,0 +1,25 @@
+from unittest import TestCase
+from pymongo import MongoClient
+
+from metrik.conf import MONGO_DATABASE, MONGO_PORT, MONGO_HOST
+from metrik.targets.mongo_target import MongoTarget
+
+
+class MongoTest(TestCase):
+    def tearDown(self):
+        super(MongoTest, self).tearDown()
+        client = MongoClient(MONGO_HOST, MONGO_PORT)
+        client.drop_database(MONGO_DATABASE)
+
+
+class MongoTestTest(MongoTest):
+    # Test that the database is dropped correctly after each test
+    def test_round1(self):
+        t = MongoTarget('consistent_collection', 'consistent_id')
+        assert not t.exists()
+        t.persist({'a': 'b'})
+
+    def test_round2(self):
+        t = MongoTarget('consistent_collection', 'consistent_id')
+        assert not t.exists()
+        t.persist({'a': 'b'})
\ No newline at end of file
diff --git a/test/targets/test_mongo_target.py b/test/targets/test_mongo_target.py
new file mode 100644
index 0000000..c44046c
--- /dev/null
+++ b/test/targets/test_mongo_target.py
@@ -0,0 +1,34 @@
+from pymongo import MongoClient
+from random import randint
+
+from metrik.targets.mongo_target import MongoTarget
+from metrik.conf import MONGO_DATABASE, MONGO_HOST, MONGO_PORT
+from test.mongo_test import MongoTest
+
+
+class MongoTargetTest(MongoTest):
+
+    def test_exists(self):
+        collection = 'test_collection'
+        id = 1234
+        t = MongoTarget(collection, id)
+        assert not t.exists()
+
+        t.persist({'a': 'b'})
+        assert t.exists()
+
+        db = MongoClient(MONGO_HOST, MONGO_PORT)[MONGO_DATABASE]
+        db_collection = db[collection]
+        db_collection.remove(id)
+        assert not t.exists()
+
+    def test_persist_retrieve(self):
+        collection = 'test_collection'
+        id = 1234
+        t = MongoTarget(collection, id)
+
+        d = {str(k): randint(0, 9999) for k in range(20)}
+        t.persist(d)
+
+        u = MongoTarget(collection, id)
+        assert u.retrieve() == d
\ No newline at end of file