add tests of input output module

ogzaf/input_output.py

@@ -17,13 +17,13 @@
 )
 
 
-def get_alpha_c(SAM=SAM, cons_dict=CONS_DICT):
+def get_alpha_c(sam=SAM, cons_dict=CONS_DICT):
     """
     Calibrate the alpha_c vector, showing the shares of household
     expenditures for each consumption category
 
     Args:
-        SAM (pd.DataFrame): SAM file
+        sam (pd.DataFrame): SAM file
         cons_dict (dict): Dictionary of consumption categories
 
     Returns:
@@ -34,8 +34,8 @@ def get_alpha_c(SAM=SAM, cons_dict=CONS_DICT):
     for key, value in cons_dict.items():
         # note the subtraction of the row to focus on domestic consumption
         category_total = (
-            SAM.loc[SAM.index.isin(value), "total"].sum()
-            - SAM.loc[SAM.index.isin(value), "row"].sum()
+            sam.loc[sam.index.isin(value), "total"].sum()
+            - sam.loc[sam.index.isin(value), "row"].sum()
         )
         alpha_c[key] = category_total
         overall_sum += category_total
@@ -45,13 +45,13 @@ def get_alpha_c(SAM=SAM, cons_dict=CONS_DICT):
     return alpha_c
 
 
-def get_io_matrix(SAM=SAM, cons_dict=CONS_DICT, prod_dict=PROD_DICT):
+def get_io_matrix(sam=SAM, cons_dict=CONS_DICT, prod_dict=PROD_DICT):
     """
     Calibrate the io_matrix array.  This array relates the share of each
     production category in each consumption category
 
     Args:
-        SAM (pd.DataFrame): SAM file
+        sam (pd.DataFrame): SAM file
         cons_dict (dict): Dictionary of consumption categories
         prod_dict (dict): Dictionary of production categories
 
@@ -71,10 +71,10 @@ def get_io_matrix(SAM=SAM, cons_dict=CONS_DICT, prod_dict=PROD_DICT):
     # the production categories from columns
     for ck, cv in cons_dict.items():
         for pk, pv in prod_dict.items():
-            io_df.loc[io_df.index == ck, pk] = SAM.loc[
-                SAM.index.isin(cv), pv
+            io_df.loc[io_df.index == ck, pk] = sam.loc[
+                sam.index.isin(cv), pv
             ].values.sum()
-    # change from levesl to share (where each row sums to one)
+    # change from levels to share (where each row sums to one)
     io_df = io_df.div(io_df.sum(axis=1), axis=0)
 
     return io_df

tests/test_input_output.py

@@ -0,0 +1,66 @@
+"""
+Tests of input_output.py module
+"""
+
+import pandas as pd
+import pytest
+from ogzaf import input_output as io
+
+
+sam_dict = {
+    # "index": ["Beer", "Chocolate", "Car", "House"],
+    "Ag": [30, 160, 0, 5],
+    "Mining": [10, 0, 100, 100],
+    "Manufacturing": [60, 40, 200, 295],
+    "total": [100, 200, 300, 400],
+    "row": [10, 20, 30, 40],
+}
+sam_df = pd.DataFrame(sam_dict, index=["Beer", "Chocolate", "Car", "House"])
+
+cons_dict = {"Food": ["Beer", "Chocolate"], "Non-food": ["Car", "House"]}
+
+prod_dict = {
+    "Primary": ["Ag", "Mining"],
+    "Secondary": ["Manufacturing"],
+}
+
+
+@pytest.mark.parametrize(
+    "sam_df, cons_dict",
+    [
+        (sam_df, cons_dict),
+    ],
+    ids=["Test 1"],
+)
+def test_get_alpha_c(sam_df, cons_dict):
+    """
+    Test of get_alpha_c() function
+    """
+    test_dict = io.get_alpha_c(sam=sam_df, cons_dict=cons_dict)
+
+    assert isinstance(test_dict, dict)
+    assert list(test_dict.keys()).sort() == ["Food", "Non-food"].sort()
+    assert test_dict["Food"] == 270 / 900
+    assert test_dict["Non-food"] == 630 / 900
+
+
+@pytest.mark.parametrize(
+    "sam_df, cons_dict, prod_dict",
+    [
+        (sam_df, cons_dict, prod_dict),
+    ],
+    ids=["Test 1"],
+)
+def test_get_io_matrix(sam_df, cons_dict, prod_dict):
+    """
+    Test of get_io_matrix()
+    """
+    test_df = io.get_io_matrix(
+        sam=sam_df, cons_dict=cons_dict, prod_dict=prod_dict
+    )
+
+    assert isinstance(test_df, pd.DataFrame)
+    assert list(test_df.columns).sort() == ["Primary", "Secondary"].sort()
+    assert list(test_df.index).sort() == ["Food", "Non-food"].sort()
+    assert test_df.loc["Food", "Primary"] == 2 / 3
+    assert test_df.loc["Food", "Secondary"] == 1 / 3