Refatorando para conseguir testar o banco de dados.

Testando e resolvendo a #26.

Author: iurisilvio

PostmonServer.py

@@ -1,7 +1,7 @@
-from bottle import route, run, error, response
+from bottle import route, run, response
 from CepTracker import CepTracker
 
-import pymongo, json, re
+from database import MongoDb as Database
 
 def expired(record_date):
 	from datetime import datetime, timedelta
@@ -22,23 +22,22 @@ def _get_info_from_correios(cep):
 
 	return info
 
+
 @route('/cep/<cep:re:\d{5}-?\d{3}>')
 def verifica_cep(cep):
 	cep = cep.replace('-','')
+	db = Database()
 
 	response.headers['Access-Control-Allow-Origin'] = '*'
-	
+
 	try:
-		con = pymongo.MongoClient('localhost')
-		db = con.postmon
-		ceps = db.ceps
-		result = ceps.find_one({'cep':cep}, fields={'_id':False})
+		result = db.get_one(cep, fields={ '_id': False })
 
 		if not result or not result.has_key('v_date') or expired(result):
 			for item in _get_info_from_correios(cep):
-				ceps.update({'cep': item['cep']}, {'$set': item}, upsert=True)
+				db.insert_or_update(item)
 
-		result = ceps.find_one({'cep':cep}, fields={'_id':False,'v_date':False})
+		result = db.get_one(cep, fields={ '_id': False, 'v_date': False })
 
 		response.headers['Cache-Control'] = 'public, max-age=2592000'
 

database.py

@@ -0,0 +1,28 @@
+import pymongo
+
+class MongoDb(object):
+
+	_fields = [
+		'logradouro',
+		'bairro',
+		'cidade',
+		'estado'
+	]
+
+	def __init__(self, address='localhost'):
+		self._client = pymongo.MongoClient(address)
+		self._db = self._client.postmon
+
+	def get_one(self, cep, **kwargs):
+		return self._db.ceps.find_one({ 'cep': cep }, **kwargs)
+
+	def insert_or_update(self, obj, **kwargs):
+
+		update = { '$set': obj }
+		empty_fields = set(self._fields) - set(obj)
+		update['$unset'] = dict((x, 1) for x in empty_fields)
+
+		self._db.ceps.update({ 'cep': obj['cep'] }, update, upsert=True)
+
+	def remove(self, cep):
+		self._db.ceps.remove({ 'cep': cep })

test/database_test.py

@@ -0,0 +1,39 @@
+# encoding: utf-8
+import unittest
+
+from database import MongoDb
+
+class MongoDbTest(unittest.TestCase):
+
+	def setUp(self):
+		self.db = MongoDb()
+
+		self.db.insert_or_update({
+			'cep': 'UNIQUE_KEY',
+			'logradouro': 'A',
+			'bairro': 'A',
+			'cidade': 'A',
+			'estado': 'A'
+		})
+
+	def test_remove_empty_fields(self):
+
+		'''
+		Quando um registro é atualizado no banco de dados,
+		as chaves inexistentes devem ser removidas.
+		'''
+
+		self.db.insert_or_update({
+			'cep': 'UNIQUE_KEY',
+			'estado': 'B'
+		})
+
+		result = self.db.get_one('UNIQUE_KEY')
+
+		self.assertEqual(result['estado'], 'B')
+		self.assertNotIn('logradouro', result)
+		self.assertNotIn('bairro', result)
+		self.assertNotIn('cidade', result)
+
+	def tearDown(self):
+		self.db.remove('UNIQUE_KEY')