diff --git a/python/tests/beagle_numba.py b/python/tests/beagle_numba.py
index 15db7ce786..8a467de828 100644
--- a/python/tests/beagle_numba.py
+++ b/python/tests/beagle_numba.py
@@ -292,7 +292,7 @@ def get_weights(typed_pos, untyped_pos, typed_cm, untyped_cm):
 
 
 @njit
-def interpolate_allele_prob(sm, ref_h, typed_pos, untyped_pos, typed_cm, untyped_cm):
+def interpolate_allele_probs(sm, ref_h, typed_pos, untyped_pos, typed_cm, untyped_cm):
     """
     Interpolate allele probabilities at the ungenotyped markers of a query haplotype
     following Equation 1 of BB2016.
@@ -410,7 +410,7 @@ def run_interpolation_beagle(
     bm = compute_backward_matrix(ref_h_typed, query_h_typed, rho, mu)
     sm = compute_state_prob_matrix(fm, bm)
     # Perform linear interpolation.
-    int_allele_prob = interpolate_allele_prob(
+    int_allele_probs = interpolate_allele_probs(
         sm,
         ref_h_untyped,
         typed_pos,
@@ -419,8 +419,8 @@ def run_interpolation_beagle(
         untyped_cm,
     )
     # Get MAP alleles at ungenotyped markers.
-    imputed_alleles, max_allele_prob = get_map_alleles(int_allele_prob)
-    return (imputed_alleles, max_allele_prob)
+    imputed_alleles, max_allele_probs = get_map_alleles(int_allele_probs)
+    return (imputed_alleles, max_allele_probs)
 
 
 def run_tsimpute(
@@ -484,7 +484,7 @@ def run_tsimpute(
     ls_hmm.backward_matrix(query_h_typed.T, fm.normalisation_factor, bm)
     sm = compute_state_prob_matrix(fm.decode(), bm.decode())
     # Perform linear interpolation.
-    int_allele_prob = interpolate_allele_prob(
+    int_allele_probs = interpolate_allele_probs(
         sm,
         ref_h_untyped,
         typed_pos,
@@ -492,5 +492,5 @@ def run_tsimpute(
         typed_cm,
         untyped_cm,
     )
-    imputed_alleles, max_allele_prob = get_map_alleles(int_allele_prob)
-    return (imputed_alleles, max_allele_prob)
+    imputed_alleles, max_allele_probs = get_map_alleles(int_allele_probs)
+    return (imputed_alleles, max_allele_probs)