Add keys strategy.

Add keys strategy, and use it.

2 files changed, 9 insertions, 2 deletions

diff --git a/tests/helpers/strategies.py b/tests/helpers/strategies.py
index de8239c..ae35d77 100644
--- a/tests/helpers/strategies.py
+++ b/tests/helpers/strategies.py
@@ -18,6 +18,7 @@
 
 from hypothesis import strategies as st
 from hypothesis.extra.pandas import column, columns, data_frames
+from scipy.stats import special_ortho_group
 
 from pyhegp.serialization import Summary, is_genotype_metadata_column, is_phenotype_metadata_column
 from pyhegp.utils import negate
@@ -88,3 +89,10 @@ def phenotype_frames(draw):
         filter(is_phenotype_metadata_column,
                phenotype.columns)),
                                      ignore_index=True)
+
+@st.composite
+def keys(draw, size):
+    return (special_ortho_group(draw(size),
+                                seed=draw(st.integers(min_value=0,
+                                                      max_value=2**32-1)))
+            .rvs())
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
index ab354e1..32cb58f 100644
--- a/tests/test_serialization.py
+++ b/tests/test_serialization.py
@@ -80,8 +80,7 @@ def test_read_write_phenotype_are_inverses(phenotype):
         file.seek(0)
         pd.testing.assert_frame_equal(phenotype, read_phenotype(file))
 
-@given(arrays("float64",
-              array_shapes(min_dims=2, max_dims=2)))
+@given(keys(st.integers(min_value=2, max_value=10)))
 def test_read_write_key_are_inverses(key):
     with tempfile.TemporaryFile() as file:
         write_key(file, key)