89 lines
3.1 KiB
Python
89 lines
3.1 KiB
Python
import unittest
|
|
from TestHelper import TestHelper
|
|
from pandas.testing import assert_frame_equal
|
|
import pandas as pd
|
|
from CountryColumnAdder import CountryColumnAdder
|
|
|
|
|
|
class CountryColumnAdderTest(unittest.TestCase):
|
|
|
|
def test_addCountryColumn(self):
|
|
# Given
|
|
dataFrame = TestHelper.createDataFrame(
|
|
columns = ['SPLTTYPE'],
|
|
data = [ ['GBPFIZER INC2020486806'],
|
|
['FRMODERNATX, INC.MOD20224'],
|
|
['dummy']],
|
|
index = pd.Index(
|
|
name = 'VAERS_ID',
|
|
data = [
|
|
"4711",
|
|
"0815",
|
|
"123"]))
|
|
countryColumnAdder = CountryColumnAdder(dataFrame)
|
|
|
|
# When
|
|
dataFrameWithCountryColumn = countryColumnAdder.addCountryColumn(dataFrame)
|
|
|
|
# Then
|
|
assert_frame_equal(
|
|
dataFrameWithCountryColumn,
|
|
TestHelper.createDataFrame(
|
|
columns = ['SPLTTYPE', 'COUNTRY'],
|
|
data = [ ['GBPFIZER INC2020486806', 'United Kingdom'],
|
|
['FRMODERNATX, INC.MOD20224', 'France'],
|
|
['dummy', None]],
|
|
index = pd.Index(
|
|
name = 'VAERS_ID',
|
|
data = [
|
|
"4711",
|
|
"0815",
|
|
"123"]),
|
|
dtypes = {'COUNTRY': 'string'}))
|
|
|
|
|
|
def test_addCountryColumn2(self):
|
|
# Given
|
|
countryColumnAdder = CountryColumnAdder(
|
|
TestHelper.createDataFrame(
|
|
columns = ['SPLTTYPE'],
|
|
data = [ ['GBPFIZER INC2020486806'],
|
|
['FRMODERNATX, INC.MOD20224'],
|
|
['dummy']],
|
|
index = pd.Index(
|
|
name = 'VAERS_ID',
|
|
data = [
|
|
2547744,
|
|
2547730,
|
|
2540815])))
|
|
dataFrame = TestHelper.createDataFrame(
|
|
columns = ['VAX_LOT'],
|
|
data = [ ['1808982'],
|
|
['EW0175'],
|
|
['EW0176']],
|
|
index = pd.Index(
|
|
name = 'VAERS_ID',
|
|
data = [
|
|
2547730,
|
|
2547730,
|
|
2547744]))
|
|
|
|
# When
|
|
dataFrameWithCountryColumn = countryColumnAdder.addCountryColumn(dataFrame)
|
|
|
|
# Then
|
|
assert_frame_equal(
|
|
dataFrameWithCountryColumn,
|
|
TestHelper.createDataFrame(
|
|
columns = ['VAX_LOT', 'COUNTRY'],
|
|
data = [ ['1808982', 'France'],
|
|
['EW0175', 'France'],
|
|
['EW0176', 'United Kingdom']],
|
|
index = pd.Index(
|
|
name = 'VAERS_ID',
|
|
data = [
|
|
2547730,
|
|
2547730,
|
|
2547744]),
|
|
dtypes = {'COUNTRY': 'string'}))
|