SS 2025

2025_Research_Software_Engineering_2/tests/test_collapse_columns.py

@@ -0,0 +1,53 @@
+"""
+Unit tests for the collapse_columns function, which merges multiple columns
+by summing their values.
+"""
+from unittest.mock import patch
+import pandas as pd
+from collapse_columns import collapse_columns
+
+
+@patch("utils.save_df")
+@patch("utils.read_df")
+@patch("checks.check_dir")
+@patch("checks.check_path")
+def test_sum_values_correctly(_mock_check_path, _mock_check_dir,
+                              mock_read_df, mock_save_df):
+    """
+    Test that collapse_columns sums values across all columns and stores the
+    result in the specified column.
+    """
+    input_df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
+    column_name = "sum"
+
+    expected_df = pd.DataFrame({column_name: [5, 7, 9]})
+    mock_read_df.return_value = input_df
+    collapse_columns("in.csv", "out.csv", column_name)
+    result_df = mock_save_df.call_args[0][0]
+
+    pd.testing.assert_frame_equal(result_df, expected_df)
+
+
+@patch("utils.save_df")
+@patch("utils.read_df")
+@patch("checks.check_dir")
+@patch("checks.check_path")
+def test_output_dimension(_mock_check_path,
+                          _mock_check_dir, mock_read_df, mock_save_df):
+    """
+    Test that the output DataFrame has one column after merging and retains
+    the original number of rows.
+    """
+    input_df = pd.DataFrame({
+        "X": [1, 2],
+        "Y": [3, 4]
+    })
+    column_name = "sum"
+
+    mock_read_df.return_value = input_df
+    expected_rows = input_df.shape[0]
+    collapse_columns("in.csv", "out.csv", column_name)
+    result_df = mock_save_df.call_args[0][0]
+
+    assert result_df.shape[1] == 1, "Output should have exactly one column"
+    assert result_df.shape[0] == expected_rows, "Output should have same number of rows as input"