continuing

src/CountriesColumnAdderTest.py

@@ -0,0 +1,31 @@
+import unittest
+from TestHelper import TestHelper
+from pandas.testing import assert_frame_equal
+import pandas as pd
+from CountriesColumnAdder import CountriesColumnAdder
+
+class CountriesColumnAdderTest(unittest.TestCase):
+    
+    def test_addCountriesColumn(self):
+        # Given
+        countriesByBatchcodeTable = TestHelper.createDataFrame(
+            columns = ['United States', 'Germany', 'Italy'],
+            data = [  [20,              0,         3]],
+            index = pd.Index(
+                    name = 'Batchcode',
+                    data = ['FE6208']))
+        columnName = 'Countries guessed'
+
+        # When
+        countriesByBatchcodeTableWithCountriesColumn = CountriesColumnAdder().addCountriesColumn(countriesByBatchcodeTable, columnName)
+
+        # Then
+        assert_frame_equal(
+            countriesByBatchcodeTableWithCountriesColumn,
+            TestHelper.createDataFrame(
+                columns = ['United States', 'Germany', 'Italy', columnName],
+                data = [  [20,              0,         3,       {'United States', 'Italy'}]],
+                index = pd.Index(
+                        name = 'Batchcode',
+                        data = ['FE6208'])))
+