Fix tests

src/TreatmentPanels.jl

@@ -3,6 +3,17 @@ module TreatmentPanels
 include("TreatmentPanel.jl")
 include("show_and_plot.jl")
 
+# Export main types
+export TreatmentPanel
 export BalancedPanel, UnbalancedPanel
 
+# Export treatment description types
+export UnitTreatmentType
+export SingleUnitTreatment, MultiUnitSimultaneousTreatment, MultiUnitStaggeredTreatment
+export TreatmentDurationType 
+export ContinuousTreatment, StartEndTreatment
+
+# Export utility functions
+export treated_ids, treated_labels, first_treated_period_ids, first_treated_period_labels, length_T₀, length_T₁
+
 end

test/runtests.jl

@@ -1,18 +1,21 @@
 using TreatmentPanels
 using Test
+using DataFrames
 
 @testset "Input tests" begin
-    test_df = DataFrame(id = ['a', 'a', 'b', 'b'], period = [1, 2, 1, 2], value = 1.0:4.0)
+    test_df = DataFrame(id = ["a", "a", "b", "b"], period = [1, 2, 1, 2], value = 1.0:4.0)
 
     # Basic functionality
-    @test BalancedPanel(test_df, :id, :period, :value, ['a' => 2]) isa BalancedPanel
-    @test BalancedPanel(test_df, :id, :period, :value, 'a' => 2) isa BalancedPanel
+    @test BalancedPanel(test_df, ["a" => 2];
+        id_var = :id, t_var = :period, outcome_var = :value) isa BalancedPanel{SingleUnitTreatment, ContinuousTreatment}
+    @test BalancedPanel(test_df, "a" => 2;
+        id_var = :id, t_var = :period, outcome_var = :value) isa BalancedPanel{SingleUnitTreatment, ContinuousTreatment}
 
     # Column not present in data
-    @test_throws "" BalancedPanel(test_df, :wrong_var, :period, :value, ['a'=> 2])
-    @test_throws "" BalancedPanel(test_df, :id, :wrong_var, :value, ['a' => 2])
-    @test_throws "" BalancedPanel(test_df, :id, :period, :wrong_var, ['a' => 2])
+    @test_throws "" BalancedPanel(test_df, "a" => 2; id_var = :wrong_var, t_var = :period, outcome_var = :value)
+    @test_throws "" BalancedPanel(test_df, "a" => 2; id_var = :id, t_var = :wrong_var, outcome_var = :value)
+    @test_throws "" BalancedPanel(test_df, "a" => 2; id_var = :id, t_var = :period, outcome_var = :wrong_var)
 
     # Treatment assignment has nonexistent units/time periods
-    @test_throws "" BalancedPanel(test_df, :id, :period, :value, ['c' => 2])
+    @test_throws "" BalancedPanel(test_df, "c" => 2; id_var = :id, t_var = :period, outcome_var = :value)
 end