diff --git a/Project.toml b/Project.toml
index 7682a45..2a2477e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "FeatureTransforms"
 uuid = "8fd68953-04b8-4117-ac19-158bf6de9782"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.3.4"
+version = "0.3.5"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
diff --git a/src/test_utils.jl b/src/test_utils.jl
index ba92c27..7191ab5 100644
--- a/src/test_utils.jl
+++ b/src/test_utils.jl
@@ -38,19 +38,19 @@ for C in (:OneToOne, :OneToMany, :ManyToOne, :ManyToMany)
 end
 
 function FeatureTransforms._apply(A, ::FakeOneToOneTransform; kwargs...)
-    return ones(size(A))
+    return replace(one, A)
 end
 
 function FeatureTransforms._apply(A, ::FakeOneToManyTransform; kwargs...)
-    return hcat(ones(size(A)), ones(size(A)))
+    return hcat(replace(one, A), replace(one, A))
 end
 
-function FeatureTransforms._apply(A, ::FakeManyToOneTransform; dims, kwargs...)
-    return ones(size(first(A)))
+function FeatureTransforms._apply(A, ::FakeManyToOneTransform; kwargs...)
+    return replace(one, first(A))
 end
 
 function FeatureTransforms._apply(A, ::FakeManyToManyTransform; kwargs...)
-    return hcat(ones(size(A)), ones(size(A)))
+    return hcat(replace(one, A), replace(one, A))
 end
 
 """
diff --git a/test/test_utils.jl b/test/test_utils.jl
index 7922825..c15dd0e 100644
--- a/test/test_utils.jl
+++ b/test/test_utils.jl
@@ -11,6 +11,13 @@ using FeatureTransforms.TestUtils
 
         M = reshape(1:9, 3, 3)
         @test FeatureTransforms.apply(M, t) == ones(3, 3)
+
+        # TODO: remove when refactoring tests
+        M = KeyedArray(M; a=["a", "b", "c"], b=[:x, :y, :z])
+        expected = KeyedArray(ones(3, 3); a=["a", "b", "c"], b=[:x, :y, :z])
+        result = FeatureTransforms.apply(M, t)
+        @test result == expected
+        @test result isa KeyedArray
     end
 
     @testset "FakeOneToManyTransform" begin
@@ -22,6 +29,13 @@ using FeatureTransforms.TestUtils
 
         M = reshape(1:9, 3, 3)
         @test FeatureTransforms.apply(M, t) == ones(3, 6)
+
+        # TODO: remove when refactoring tests
+        M = KeyedArray(M; a=["a", "b", "c"], b=[:x, :y, :z])
+        expected = KeyedArray(ones(3, 6); a=["a", "b", "c"], b=[:x, :y, :z, :x, :y, :z])
+        result = FeatureTransforms.apply(M, t)
+        @test result == expected
+        @test result isa KeyedArray
     end
 
     @testset "FakeManyToOneTransform" begin
@@ -33,6 +47,13 @@ using FeatureTransforms.TestUtils
 
         M = reshape(1:9, 3, 3)
         @test FeatureTransforms.apply(M, t; dims=1) == ones(3)
+
+        # TODO: remove when refactoring tests
+        M = KeyedArray(M; a=["a", "b", "c"], b=[:x, :y, :z])
+        expected = KeyedArray(ones(3); a=["a", "b", "c"])
+        result = FeatureTransforms.apply(M, t; dims=:b)
+        @test result == expected
+        @test result isa KeyedArray
     end
 
     @testset "FakeManyToManyTransform" begin
@@ -44,6 +65,13 @@ using FeatureTransforms.TestUtils
 
         M = reshape(1:9, 3, 3)
         @test FeatureTransforms.apply(M, t) == ones(3, 6)
+
+        # TODO: remove when refactoring tests
+        M = KeyedArray(M; a=["a", "b", "c"], b=[:x, :y, :z])
+        expected = KeyedArray(ones(3, 6); a=["a", "b", "c"], b=[:x, :y, :z, :x, :y, :z])
+        result = FeatureTransforms.apply(M, t)
+        @test result == expected
+        @test result isa KeyedArray
     end