ajout u12

MTES-MCT · Jan 21, 2025 · 335e798 · 335e798
1 parent 94af845
commit 335e798
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/src/prefect/indicators/usage/u12.py b/src/prefect/indicators/usage/u12.py
@@ -25,12 +25,16 @@
 )
 
 NUM_POC_IN_OPERATION_FOR_LEVEL_QUERY_TEMPLATE = """
+        WITH
+            $power_range
         SELECT
             count(*) AS value,
+            category,
             level_id
         FROM (
             SELECT 
                 point_de_charge_id,
+                category,
                 $level_id AS level_id
             FROM
                 SESSION
@@ -44,10 +48,12 @@
                 AND $level_id IN ($indexes)
             GROUP BY 
                 point_de_charge_id,
+                category,
                 $level_id
             ) AS liste_pdc
         GROUP BY 
-            level_id
+            level_id,
+            category
         """
 
 QUERY_NATIONAL_TEMPLATE = """
@@ -120,7 +126,7 @@ def u12_for_level(
         "level": level,
         "period": timespan.period,
         "timestamp": timespan.start.isoformat(),
-        "category": None,
+        "category": merged["category"].astype("str"),
         "extras": None,
     }
     return pd.DataFrame(indicators)
@@ -144,7 +150,7 @@ def u12_national(timespan: IndicatorTimeSpan) -> pd.DataFrame:
         "level": Level.NATIONAL,
         "period": timespan.period,
         "timestamp": timespan.start.isoformat(),
-        "category": None,
+        "category": res["category"].astype("str"),,
         "extras": None,
     }
     return pd.DataFrame(indicators)

diff --git a/src/prefect/tests/usage/test_u12.py b/src/prefect/tests/usage/test_u12.py
@@ -121,6 +121,11 @@ def test_flow_u12_calculate(db_connection):
     )
     expected = sum(result.one()) + 1
     all_levels = [Level.NATIONAL, Level.REGION, Level.DEPARTMENT, Level.CITY, Level.EPCI]
+    indicators = test_flow_u12_for_level_with_various_chunk_sizes.calculate(TIMESPAN, all_levels, create_artifact=True, format_pd=True)
+    assert indicators["value"].sum() == expected
+
+
+
     indicators = u12.calculate(TIMESPAN, all_levels, create_artifact=True)
     assert len(indicators) == expected