Merge 83dde70 into 196dd67

rosie/chamber_of_deputies/tests/test_irregular_companies_classifier.py

@@ -1,70 +1,61 @@
-import datetime
+from collections import namedtuple
+from datetime import date
+from itertools import chain
 from unittest import TestCase
 
 import pandas as pd
 
-from rosie.chamber_of_deputies.classifiers.irregular_companies_classifier import IrregularCompaniesClassifier
+from rosie.chamber_of_deputies.classifiers import IrregularCompaniesClassifier
+
+
+Status = namedtuple('Status', ('situation_date', 'issue_date', 'expected'))
 
 
 class TestIrregularCompaniesClassifier(TestCase):
 
+    SUSPICIOUS_SITUATIONS = (
+        'BAIXADA',
+        'NULA',
+        'INAPTA',
+        'SUSPENSA',
+    )
+
+    SITUATIONS = chain(SUSPICIOUS_SITUATIONS, ('ABERTA',))
+
+    STATUS = (
+        Status(date(2013, 1, 30), date(2013, 1, 1), False),
+        Status(date(2013, 1, 1), date(2013, 1, 30), True)
+    )
+
     def setUp(self):
         self.subject = IrregularCompaniesClassifier()
 
-    def _get_company_dataset(self, fields):
+    def _get_company_dataset(self, **kwargs):
         base_company = {
             'recipient_id': '02989654001197',
-            'situation_date': datetime.date(2013, 1, 3),
-            'issue_date': datetime.date(2013, 1, 30),
+            'situation_date': date(2013, 1, 3),
+            'issue_date': date(2013, 1, 30),
             'situation': '',
         }
-        for key, value in fields.items():
-            base_company[key] = value
-        dataset = pd.DataFrame([base_company, ])
+        base_company.update(kwargs)
+        dataset = pd.DataFrame([base_company])
         return dataset
 
     def test_is_regular_company(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'ABERTA'
-            }))[0], False)
-
-    def test_is_irregular_company_BAIXADA(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'BAIXADA'
-            }))[0], True)
-
-    def test_is_irregular_company_NULA(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'NULA'
-            }))[0], True)
-
-    def test_is_irregular_company_INAPTA(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'INAPTA'
-            }))[0], True)
-
-    def test_is_irregular_company_SUSPENSA(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'SUSPENSA'
-            }))[0], True)
-
-    def test_is_valid_if_suspended_after_expense(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'SUSPENSA',
-                'situation_date': datetime.date(2013, 1, 30),
-                'issue_date': datetime.date(2013, 1, 1),
-            }))[0], False)
-
-    def test_is_irregular_if_suspended_before_expense(self):
-        self.assertEqual(
-            self.subject.predict(self._get_company_dataset({
-                'situation': 'SUSPENSA',
-                'situation_date': datetime.date(2013, 1, 1),
-                'issue_date': datetime.date(2013, 1, 30),
-            }))[0], True)
+        for situation in self.SITUATIONS:
+            company = self._get_company_dataset(situation=situation)
+            expected = situation in self.SUSPICIOUS_SITUATIONS
+            result, *_ = self.subject.predict(company)
+            with self.subTest():
+                self.assertEqual(result, expected, msg=company)
+
+    def test_if_company_is_suspended(self):
+        for status in self.STATUS:
+            company = self._get_company_dataset(
+                situation='SUSPENSA',
+                situation_date=status.situation_date,
+                issue_date=status.issue_date,
+            )
+            result, *_ = self.subject.predict(company)
+            with self.subTest():
+                self.assertEqual(result, status.expected, msg=company)