Resolvendo o bug #16

PostmonServer.py

-from bottle import route, run, error, response
+from bottle import route, run, response
 from CepTracker import CepTracker
 from requests import ConnectionError
 
-import pymongo, json, re
+from database import MongoDb as Database
 
 def expired(record_date):
 	from datetime import datetime, timedelta
@@ -23,27 +23,26 @@ 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):
 			try:
 				for item in _get_info_from_correios(cep):
-					ceps.update({'cep': item['cep']}, {'$set': item}, upsert=True)
+					db.insert_or_update(item)
 
 			except ConnectionError:
 				response.status = '503 Servico Temporariamente Indisponivel'
 
-		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'
 

build/pip-delete-this-directory.txt

+This file is placed here by pip to indicate the source was put
+here by pip.
+
+Once this package is successfully installed this source code will be
+deleted (unless you remove this file).

database.py

+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

+# 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')

test/postmon_test.py

@@ -107,6 +107,7 @@ class PostmonWebTest(unittest.TestCase, PostmonBaseTest):
 		except webtest.AppError as ex:
 			if not expected and '404' in ex.message and cep in ex.message:
 				return
+			raise ex
 
 		for k, v in expected[0].items():
 			self.assertEqual(v, result[k])