Added a few AxisArray tests.

Project.toml

@@ -4,7 +4,6 @@ authors = ["Invenia Technical Computing"]
 version = "0.2.0"
 
 [deps]
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 Missings = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -19,9 +18,11 @@ Tables = "0.2"
 julia = "1"
 
 [extras]
+AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 RDatasets = "ce6b1742-4840-55fa-b093-852dadbb1d8b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DataFrames", "RDatasets", "Test"]
+test = ["AxisArrays", "DataFrames", "Dates", "RDatasets", "Test"]

test/runtests.jl

@@ -1,7 +1,9 @@
 using Impute
 using Tables
 using Test
+using AxisArrays
 using DataFrames
+using Dates
 using RDatasets
 using Statistics
 using StatsBase
@@ -29,6 +31,12 @@ import Impute:
     # We call collect to not have a wrapper type that references the same data.
     m = collect(reshape(a, 5, 4))
 
+    aa = AxisArray(
+        deepcopy(m),
+        Axis{:time}(DateTime(2017, 6, 5, 5):Hour(1):DateTime(2017, 6, 5, 9)),
+        Axis{:id}(1:4)
+    )
+
     table = DataFrame(
         :sin => allowmissing(sin.(1.0:1.0:20.0)),
         :cos => allowmissing(sin.(1.0:1.0:20.0)),
@@ -111,8 +119,6 @@ import Impute:
                     result = impute(rowtab, DropObs(; context=ctx))
                     expected = Tables.rowtable(dropmissing(table))
 
-                    @show result
-                    @show expected
                     @test isequal(result, expected)
                     @test isequal(result, Impute.dropobs(rowtab; context=ctx))
 
@@ -123,6 +129,22 @@ import Impute:
                     @test isequal(rowtab_, expected)
                 end
             end
+
+            @testset "AxisArray" begin
+                # Because we're removing 2 of our 5 rows we need to change the limit.
+                ctx = Context(; limit=0.4)
+                result = impute(aa, DropObs(; context=ctx))
+                expected = aa[[1, 4, 5], :]
+
+                @test isequal(result, expected)
+                @test isequal(result, Impute.dropobs(aa; context=ctx))
+
+                aa_ = Impute.dropobs!(aa; context=ctx)
+                # The mutating test is broken because we need to making a copy of
+                # the original matrix
+                @test_broken isequal(aa, expected)
+                @test isequal(aa_, expected)
+            end
         end
 
         @testset "DropVars" begin
@@ -152,8 +174,6 @@ import Impute:
                     result = impute(df, DropVars(; context=ctx))
                     expected = select(df, :cos)
 
-                    @show result
-                    @show expected
                     @test isequal(result, expected)
                     @test isequal(result, Impute.dropvars(df; context=ctx))
 
@@ -192,6 +212,20 @@ import Impute:
                     @test_broken isequal(rowtab, expected)
                 end
             end
+            @testset "AxisArray" begin
+                ctx = Context(; limit=0.5)
+                result = impute(aa, DropVars(; context=ctx))
+                expected = copy(aa)[:, 3:4]
+
+                @test isequal(result, expected)
+                @test isequal(result, Impute.dropvars(aa; context=ctx))
+
+                aa_ = Impute.dropvars!(aa; context=ctx)
+                # The mutating test is broken because we need to making a copy of
+                # the original matrix
+                @test_broken isequal(aa, expected)
+                @test isequal(aa_, expected)
+            end
         end
     end
 
@@ -383,6 +417,19 @@ import Impute:
             # Confirm that we don't have any more missing values
             @test all(!ismissing, result)
         end
+
+        @testset "AxisArray" begin
+            data = AxisArray(
+                Matrix(orig),
+                Axis{:row}(1:size(orig, 1)),
+                Axis{:V}(names(orig)),
+            )
+            result = Impute.interp(data; context=ctx) |> Impute.locf!() |> Impute.nocb!()
+
+            @test size(result) == size(data)
+            # Confirm that we don't have any more missing values
+            @test all(!ismissing, result)
+        end
     end
 
     @testset "Alternate missing functions" begin