diff --git a/python/tests/test_imputation.py b/python/tests/test_imputation.py
index e87dc5b6a6..0a5c0f656e 100644
--- a/python/tests/test_imputation.py
+++ b/python/tests/test_imputation.py
@@ -239,13 +239,11 @@ def get_test_data(matrix_text, par):
         # 0 = ref allele, 1 = alt allele
         return x[:, 7].reshape((4, 4))[:, 0]
     elif par == "shift":
-        # Shift factor
-        # TODO
-        pass
+        # Shift factor, one per site
+        return x[:, 8].reshape((4, 4))[:, 0]
     elif par == "scale":
-        # Scale factor
-        # TODO
-        pass
+        # Scale factor, one per site
+        return x[:, 9].reshape((4, 4))[:, 0]
     elif par == "sum":
         # Sum of values over haplotypes
         return x[:, 10].reshape((4, 4))[:, 0]