diff --git a/python/tests/test_beagle.py b/python/tests/test_beagle.py
index 5e410665c9..c42c889598 100644
--- a/python/tests/test_beagle.py
+++ b/python/tests/test_beagle.py
@@ -461,25 +461,23 @@ def _compute_allele_probabilities(a):
     return p
 
 
-def get_map_alleles(i_hap_probs):
+def get_map_alleles(allele_probs):
     """
     Compute the maximum a posteriori alleles at the imputed markers of a query haplotype.
 
     Assuming all biallelic sites, it is a matrix of size (x, 2),
     where x is the number of imputed markers.
 
-    This implementation is based on Equation 2 (the rearranged one) in BB2016.
-
-    :param numpy.ndarray i_hap_probs: Interpolated haplotype probability matrix.
-    :return: MAP alleles at imputed markers.
+    :param numpy.ndarray allele_probs: Interpolated allele probabilities.
+    :return: Imputed alleles in the query haplotype.
     :rtype: numpy.ndarray
     """
-    x = len(i_hap_probs)
+    x = len(allele_probs)
     imputed_alleles = np.zeros((x, 2))
     # TODO: Vectorise over the imputed markers
     for i in np.arange(x):
-        imputed_alleles[i, 0] = np.argmax(i_hap_probs[i, :])
-        imputed_alleles[i, 1] = np.argmax(i_hap_probs[i, :])
+        imputed_alleles[i, 0] = np.argmax(allele_probs[i, :])
+        imputed_alleles[i, 1] = np.argmax(allele_probs[i, :])
     return imputed_alleles