Merge pull request #1 from ihmeuw-msca/bugfix/tests

bugfix/tests

tests/conftest.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+EXAMPLES = Path(__file__).resolve().parent / "examples"
+
+
+class Example1D:
+    def __init__(self):
+        self.df_obs = pd.read_csv(EXAMPLES / "example_1D" / "observations.csv")
+        self.df_margin = pd.read_csv(EXAMPLES / "example_1D" / "margin.csv")
+
+
+class Example2D:
+    def __init__(self):
+        self.df_obs = pd.read_csv(EXAMPLES / "example_2D" / "observations.csv")
+        self.df_margins_1 = pd.read_csv(
+            EXAMPLES / "example_2D" / "margins_1.csv"
+        )
+        self.df_margins_2 = pd.read_csv(
+            EXAMPLES / "example_2D" / "margins_2.csv"
+        )
+
+
+class Example3D:
+    def __init__(self):
+        self.df_obs = pd.read_csv(EXAMPLES / "example_3D" / "observations.csv")
+        self.df_margins_1 = pd.read_csv(
+            EXAMPLES / "example_3D" / "margins_1.csv"
+        )
+        self.df_margins_2 = pd.read_csv(
+            EXAMPLES / "example_3D" / "margins_2.csv"
+        )
+        self.df_margins_3 = pd.read_csv(
+            EXAMPLES / "example_3D" / "margins_3.csv"
+        )
+
+
+@pytest.fixture
+def example_1D():
+    return Example1D()
+
+
+@pytest.fixture
+def example_2D():
+    return Example2D()
+
+
+@pytest.fixture
+def example_3D():
+    return Example3D()

tests/test_run.py

@@ -1,25 +1,35 @@
-import pytest
-import numpy as np
-import os
-import pandas as pd
-
 from raking.run_raking import run_raking
 
-def test_run_raking_1D():
-    df_obs = pd.read_csv('examples/example_1D/observations.csv')
-    df_margin = pd.read_csv('examples/example_1D/margin.csv')
-    run_raking(1, 'examples/example_1D', df_obs, [df_margin], ['var1'])
 
-def test_run_raking_2D():
-    df_obs = pd.read_csv('examples/example_2D/observations.csv')
-    df_margins_1 = pd.read_csv('examples/example_2D/margins_1.csv')
-    df_margins_2 = pd.read_csv('examples/example_2D/margins_2.csv')
-    run_raking(2, 'examples/example_2D', df_obs, [df_margins_1, df_margins_2], ['var1', 'var2'])
+def test_run_raking_1D(tmp_path, example_1D):
+    run_raking(
+        1,
+        str(tmp_path),
+        example_1D.df_obs,
+        [example_1D.df_margin],
+        ["var1"],
+    )
+
+
+def test_run_raking_2D(tmp_path, example_2D):
+    run_raking(
+        2,
+        str(tmp_path),
+        example_2D.df_obs,
+        [example_2D.df_margins_1, example_2D.df_margins_2],
+        ["var1", "var2"],
+    )
 
-def test_run_raking_3D():
-    df_obs = pd.read_csv('examples/example_3D/observations.csv')
-    df_margins_1 = pd.read_csv('examples/example_3D/margins_1.csv')
-    df_margins_2 = pd.read_csv('examples/example_3D/margins_2.csv')
-    df_margins_3 = pd.read_csv('examples/example_3D/margins_3.csv')
-    run_raking(3, 'examples/example_3D', df_obs, [df_margins_1, df_margins_2, df_margins_3], ['var1', 'var2', 'var3'])
 
+def test_run_raking_3D(tmp_path, example_3D):
+    run_raking(
+        3,
+        str(tmp_path),
+        example_3D.df_obs,
+        [
+            example_3D.df_margins_1,
+            example_3D.df_margins_2,
+            example_3D.df_margins_3,
+        ],
+        ["var1", "var2", "var3"],
+    )