From 82e1c2899d989dc3dbcdf0a44d54debfd57bc4c4 Mon Sep 17 00:00:00 2001
From: GabrielKS <23368820+GabrielKS@users.noreply.github.com>
Date: Tue, 17 Oct 2023 14:39:07 -0600
Subject: [PATCH] Fix #8, add tests using PowerSystemCaseBuilder

 - Fix #8, accomplished by the small changes in `get_data.jl`
 - Add a dependency on PowerSystemCaseBuilder
 - Test the fix
 - NOTE: depends on the new system `c_sys5_all_components` I just added
   to PowerSystemCaseBuilder, which in turn depends on a
   `convert_component!` method in PowerSystems I just added
---
 Project.toml                |  1 +
 src/get_data.jl             |  4 ++--
 test/test_result_sorting.jl | 26 +++++++++++++++++++++-----
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/Project.toml b/Project.toml
index 4ea439c..3a6b70e 100644
--- a/Project.toml
+++ b/Project.toml
@@ -10,6 +10,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 InfrastructureSystems = "2cd47ed4-ca9b-11e9-27f2-ab636a7671f1"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 PowerSimulations = "e690365d-45e2-57bb-ac84-44ba829e73c4"
+PowerSystemCaseBuilder = "f00506e0-b84f-492a-93c2-c0a9afc4364e"
 PowerSystems = "bcd98974-b02a-5e2f-9ee0-a103f5c450dd"
 TimeSeries = "9e3dc215-6440-5c97-bce1-76c03772f85e"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
diff --git a/src/get_data.jl b/src/get_data.jl
index 93d469f..ce4073b 100644
--- a/src/get_data.jl
+++ b/src/get_data.jl
@@ -425,11 +425,11 @@ get_base_power(results::PSI.ProblemResults) = results.base_power
 function get_load_data(
     system::PSY.System;
     aggregation::Union{
-        Type{PSY.StaticLoad},
+        Type{<:PSY.StaticLoad},
         Type{PSY.ACBus},
         Type{PSY.System},
         Type{<:PSY.AggregationTopology},
-    } = PSY.StandardLoad,
+    } = PSY.StaticLoad,
     kwargs...,
 )
     aggregation_components =
diff --git a/test/test_result_sorting.jl b/test/test_result_sorting.jl
index 5157bcc..7bfbcb2 100644
--- a/test/test_result_sorting.jl
+++ b/test/test_result_sorting.jl
@@ -101,16 +101,32 @@ end
 end
 
 @testset "Test system data getters" begin
-    sys = PSI.get_system(results_uc)
+    # We can test get_load_data without running a simulation
+    sys = PSB.build_system(PSB.PSITestSystems, "c_sys5_all_components")
+
     load_data1 = PA.get_load_data(sys; aggregation = ACBus)
     @test length(load_data1.data) == 3
     @test length(load_data1.time) == 24
 
-    load_data2 = PA.get_load_data(sys, aggregation = StaticLoad)
+    load_data2 = PA.get_load_data(sys; aggregation = StaticLoad)
     @test length(load_data2.data) == 3
     @test length(load_data2.time) == 24
 
-    load_data3 = PA.get_load_data(sys, aggregation = System)
-    @test length(load_data2.data) == 3
-    @test length(load_data2.time) == 24
+    # With System aggregation, everything gets lumped together
+    load_data3 = PA.get_load_data(sys; aggregation = System)
+    @test length(load_data3.data) == 1
+    @test length(load_data3.time) == 24
+
+    # No aggregation specified: default is StaticLoad
+    load_data4 = PA.get_load_data(sys)
+    @test length(load_data4.data) == 3
+    @test length(load_data4.time) == 24
+
+    load_data5 = PA.get_load_data(sys; aggregation = PowerLoad)
+    @test length(load_data5.data) == 2
+    @test length(load_data5.time) == 24
+
+    load_data6 = PA.get_load_data(sys; aggregation = StandardLoad)
+    @test length(load_data6.data) == 1
+    @test length(load_data6.time) == 24
 end