Ensure combining datasets return the same number of rows as reimburse…

…ments

rosie/dataset.py

@@ -17,6 +17,7 @@ def get(self):
         reimbursements = self.get_reimbursements()
         companies = self.get_companies()
         return pd.merge(reimbursements, companies,
+                        how='left',
                         left_on='cnpj_cpf',
                         right_on='cnpj')
 
@@ -45,7 +46,7 @@ def get_companies(self):
         is_in_brazil = ('(-73.992222 < longitude < -34.7916667) & '
                         '(-33.742222 < latitude < 5.2722222)')
         dataset = pd.read_csv(os.path.join(self.path, self.COMPANIES_DATASET),
-                              dtype={'cnpj_cpf': np.str},
+                              dtype={'cnpj': np.str},
                               low_memory=False)
         dataset = dataset.query(is_in_brazil)
         dataset['cnpj'] = dataset['cnpj'].str.replace(r'\D', '')

tests/test_dataset.py

@@ -0,0 +1,24 @@
+import os.path
+from tempfile import mkdtemp
+from unittest import TestCase
+from unittest.mock import patch
+from shutil import copy2
+
+from rosie import Dataset
+
+
+class TestDataset(TestCase):
+
+    def setUp(self):
+        temp_path = mkdtemp()
+        copy2('tests/fixtures/companies.xz',
+              os.path.join(temp_path, Dataset.COMPANIES_DATASET))
+        copy2('tests/fixtures/reimbursements.xz', temp_path)
+        self.subject = Dataset(temp_path)
+
+    @patch('rosie.dataset.CEAPDataset')
+    @patch('rosie.dataset.fetch')
+    def test_get_performs_a_left_merge_between_reimbursements_and_companies(self, _ceap_dataset, _fetch):
+        dataset = self.subject.get()
+        self.assertEqual(5, len(dataset))
+        self.assertEqual(1, dataset['legal_entity'].isnull().sum())