From bd9f3f360c765c42b4d138d8de0368131bc24586 Mon Sep 17 00:00:00 2001
From: Vincent <vincent.chery@m4x.org>
Date: Tue, 10 Dec 2024 17:13:36 +0100
Subject: [PATCH] Update extract_segments_of_year

---
 .../monitorfish/fleet_segments_of_year.sql    |  6 ++-
 .../test_shared_tasks/test_segments.py        | 40 ++++++-------------
 2 files changed, 16 insertions(+), 30 deletions(-)

diff --git a/datascience/src/pipeline/queries/monitorfish/fleet_segments_of_year.sql b/datascience/src/pipeline/queries/monitorfish/fleet_segments_of_year.sql
index cbb22e703e..444bc544d4 100644
--- a/datascience/src/pipeline/queries/monitorfish/fleet_segments_of_year.sql
+++ b/datascience/src/pipeline/queries/monitorfish/fleet_segments_of_year.sql
@@ -1,4 +1,5 @@
-SELECT 
+SELECT
+    year,
     segment,
     segment_name,
     gears,
@@ -12,4 +13,5 @@ SELECT
     vessel_types,
     impact_risk_factor
 FROM public.fleet_segments
-WHERE year = :year
\ No newline at end of file
+WHERE year = :year
+ORDER BY segment
\ No newline at end of file
diff --git a/datascience/tests/test_pipeline/test_shared_tasks/test_segments.py b/datascience/tests/test_pipeline/test_shared_tasks/test_segments.py
index 191d00130d..205b983708 100644
--- a/datascience/tests/test_pipeline/test_shared_tasks/test_segments.py
+++ b/datascience/tests/test_pipeline/test_shared_tasks/test_segments.py
@@ -11,8 +11,12 @@
 
 
 @pytest.fixture
-def expected_all_segments() -> pd.DataFrame:
-    current_year = datetime.utcnow().year
+def current_year() -> int:
+    return datetime.utcnow().year
+
+
+@pytest.fixture
+def expected_all_segments(current_year) -> pd.DataFrame:
     expected_segments = pd.DataFrame(
         {
             "year": [2022, 2022, current_year, current_year],
@@ -59,38 +63,18 @@ def expected_all_segments() -> pd.DataFrame:
 
 
 @pytest.fixture
-def expected_segments_of_year() -> pd.DataFrame:
-    return pd.DataFrame(
-        {
-            "segment": ["SWW01/02/03", "SWW04"],
-            "segment_name": ["Bottom trawls", "Midwater trawls"],
-            "gears": [
-                ["OTB", "OTT", "PTB", "OT", "PT", "TBN", "TBS", "TX", "TB"],
-                ["OTM", "PTM"],
-            ],
-            "fao_areas": [["27.8.c", "27.8", "27.9"], ["27.8.c", "27.8"]],
-            "min_mesh": [80.0, None],
-            "max_mesh": [120.0, None],
-            "target_species": [["HKE", "SOL", "ANF", "MNZ", "NEP", "LEZ"], ["HKE"]],
-            "min_share_of_target_species": [0.0, 0.0],
-            "main_scip_species_type": ["DEMERSAL", "PELAGIC"],
-            "priority": [0.0, 1.0],
-            "vessel_types": [
-                None,
-                ["Navire qui pêche", "Chalutier", "Ligneur", "Navire qui navigue"],
-            ],
-            "impact_risk_factor": [3.0, 2.1],
-        }
-    )
+def expected_segments_of_current_year(expected_all_segments) -> pd.DataFrame:
+    return expected_all_segments.loc[2:4].reset_index(drop=True)
 
 
-def test_extract_segments_of_year(reset_test_data, expected_segments_of_year):
-    current_year = datetime.utcnow().year
+def test_extract_segments_of_year(
+    reset_test_data, expected_segments_of_current_year, current_year
+):
     segments = extract_segments_of_year.run(current_year)
 
     pd.testing.assert_frame_equal(
         segments.sort_values("segment").reset_index(drop=True),
-        expected_segments_of_year,
+        expected_segments_of_current_year,
     )