Parameterize presence of reference column in genotype frame strategy.

1 files changed, 5 insertions, 3 deletions

diff --git a/tests/helpers/strategies.py b/tests/helpers/strategies.py
index ae35d77..ed9adde 100644
--- a/tests/helpers/strategies.py
+++ b/tests/helpers/strategies.py
@@ -60,11 +60,13 @@ def summaries(draw):
                                           elements=st.floats(allow_nan=False))))))
 
 @st.composite
-def genotype_frames(draw):
+def genotype_frames(draw,
+                    reference_present=st.booleans()):
     genotype = draw(data_frames(
         columns=([chromosome_column, position_column]
-                 + ([reference_column] if draw(st.booleans()) else [])
-                 + columns(draw(sample_names),
+                 + ([reference_column]
+                    if draw(reference_present)
+                    else [])
                            dtype="float64",
                            elements=st.floats(allow_nan=False)))))
     return genotype.drop_duplicates(subset=list(