further tests

moj-analytical-services · Jun 3, 2020 · c77f1e4 · c77f1e4
1 parent 7f18046
commit c77f1e4
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/tests/test_fix_string.py b/tests/test_fix_string.py
@@ -54,4 +54,13 @@ def test_fix_2(spark):
 
     df_expected = pd.DataFrame(df_expected)
 
-    pd.testing.assert_frame_equal(df_result,df_expected)
+    pd.testing.assert_frame_equal(df_result,df_expected)
+
+    # Check col not dropped if names are the same
+    df = spark.createDataFrame(Row(**x) for x in data_list)
+    df = remove_leading_zeros(df, "id_var", "id_var")
+    df_expected = df_expected.rename(columns={"id_var_norm": "id_var"})
+
+
+    df_result = df.toPandas()
+    pd.testing.assert_frame_equal(df_result,df_expected)