added test for correct shape when extracting

IndEcol · Apr 9, 2024 · cc968fe · cc968fe
1 parent b9959c3
commit cc968fe
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -304,6 +304,19 @@ def test_extract(fix_testmrio):
     assert "F" in with_missing.keys()
     assert "FOO" not in with_missing.keys()
 
+    # Test for correct shape when extracting one row
+    assert tt.factor_inputs.extract("Value Added", return_type="ext").F.index == tt.factor_inputs.get_rows()
+    assert tt.factor_inputs.extract(("Value Added"), return_type="ext").F.index == tt.factor_inputs.get_rows()
+    assert tt.factor_inputs.extract(["Value Added"], return_type="ext").F.index == tt.factor_inputs.get_rows()
+
+    assert tt.factor_inputs.extract(tt.factor_inputs.get_rows(), return_type="ext").F.index == tt.factor_inputs.get_rows()
+    pdt.assert_index_equal(tt.emissions.extract(tt.emissions.get_rows(), return_type="ext").F.index, tt.emissions.get_rows())
+    assert tt.emissions.extract(tt.emissions.get_rows()[0], return_type="ext").F.index == tt.emissions.get_rows()[0]
+
+
+
+
+
 
 def test_extension_extract(fix_testmrio):
     tt = fix_testmrio.testmrio
@@ -343,6 +356,7 @@ def test_extension_extract(fix_testmrio):
     assert all(dfm.get_sectors() == tt.get_sectors())
 
 
+
 def test_diag_stressor(fix_testmrio):
     stressor_name = ("emission_type1", "air")
     stressor_number = 0