diff --git a/tests/statespace/test_distributions.py b/tests/statespace/test_distributions.py
index fe9a0df1..4d0d2781 100644
--- a/tests/statespace/test_distributions.py
+++ b/tests/statespace/test_distributions.py
@@ -11,12 +11,17 @@
 from pymc_extras.statespace.filters.distributions import (
     LinearGaussianStateSpace,
     LinearGaussianStateSpaceRV,
+    SequenceMvNormal,
+    _LinearGaussianStateSpace,
 )
 from pymc_extras.statespace.utils.constants import (
+    ALL_STATE_DIM,
     JITTER_DEFAULT,
     LONG_MATRIX_NAMES,
     MISSING_FILL,
+    OBS_STATE_DIM,
     SHORT_NAME_TO_LONG,
+    TIME_DIM,
 )
 from tests.statespace.utilities.shared_fixtures import (  # pylint: disable=unused-import
     rng,