From 5b97a40f493dfd250831d303f5e3e4ba0a8a2c08 Mon Sep 17 00:00:00 2001
From: Arun Isaac
Date: Thu, 4 Sep 2025 20:42:47 +0100
Subject: Add keys strategy.

Add keys strategy, and use it.
---
 tests/helpers/strategies.py | 8 ++++++++
 tests/test_serialization.py | 3 +--
 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)
-- 
cgit 1.4.1