From 335e798eb72ce448ba9191ea571d29a6350ff7a0 Mon Sep 17 00:00:00 2001
From: philippe <philippe@loco-labs.io>
Date: Tue, 21 Jan 2025 23:22:55 +0100
Subject: [PATCH] ajout u12

---
 src/prefect/indicators/usage/u12.py | 12 +++++++++---
 src/prefect/tests/usage/test_u12.py |  5 +++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/prefect/indicators/usage/u12.py b/src/prefect/indicators/usage/u12.py
index 2db0a094..49531fc7 100644
--- 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
index 551c7eff..dde4b1ee 100644
--- 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