diff --git a/test_libs/pyspec/eth2spec/test/context.py b/test_libs/pyspec/eth2spec/test/context.py
index b2e67d761..d3d9be588 100644
--- a/test_libs/pyspec/eth2spec/test/context.py
+++ b/test_libs/pyspec/eth2spec/test/context.py
@@ -6,7 +6,7 @@ from .helpers.genesis import create_genesis_state
 from .utils import spectest, with_args
 
 # Provides a genesis state as first argument to the function decorated with this
-with_state = with_args(lambda: [create_genesis_state(spec.SLOTS_PER_EPOCH * 8)])
+with_state = with_args(lambda: [create_genesis_state(spec.SHARD_COUNT * 2)])
 
 
 # shorthand for decorating @with_state @spectest()
diff --git a/test_libs/pyspec/eth2spec/test/helpers/keys.py b/test_libs/pyspec/eth2spec/test/helpers/keys.py
index fe27b78a1..f59fbb159 100644
--- a/test_libs/pyspec/eth2spec/test/helpers/keys.py
+++ b/test_libs/pyspec/eth2spec/test/helpers/keys.py
@@ -1,5 +1,6 @@
 from py_ecc import bls
+from eth2spec.phase0 import spec
 
-privkeys = [i + 1 for i in range(1024)]
+privkeys = [i + 1 for i in range(spec.SHARD_COUNT * 8)]
 pubkeys = [bls.privtopub(privkey) for privkey in privkeys]
 pubkey_to_privkey = {pubkey: privkey for privkey, pubkey in zip(privkeys, pubkeys)}
diff --git a/test_libs/pyspec/eth2spec/test/test_finality.py b/test_libs/pyspec/eth2spec/test/test_finality.py
index 732dd8f0c..f47ee5f77 100644
--- a/test_libs/pyspec/eth2spec/test/test_finality.py
+++ b/test_libs/pyspec/eth2spec/test/test_finality.py
@@ -5,11 +5,8 @@ from eth2spec.phase0.state_transition import (
     state_transition,
 )
 from .context import spec_state_test
-from .helpers.state import (
-    next_epoch,
-    apply_empty_block
-)
-from .helpers.block import build_empty_block_for_next_slot
+from .helpers.state import next_epoch
+from .helpers.block import build_empty_block_for_next_slot, apply_empty_block
 from .helpers.attestations import (
     fill_aggregate_attestation,
     get_current_epoch,