From 048783ceee5e420510162c4417af80fff69fc20c Mon Sep 17 00:00:00 2001
From: Alvaro Tinoco <alvarotinocomarmol@gmail.com>
Date: Sun, 19 Sep 2021 09:44:23 +0000
Subject: [PATCH] feat(method): add get_consumption_data

---
 datadis/__init__.py   | 26 ++++++++++++++++++++++++--
 datadis/types.py      | 10 +++++++++-
 tests/test_datadis.py | 21 +++++++++++++++++++--
 3 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/datadis/__init__.py b/datadis/__init__.py
index 486126c..b69080d 100644
--- a/datadis/__init__.py
+++ b/datadis/__init__.py
@@ -1,5 +1,6 @@
-from typing import List
-from datadis.types import ContractDetail, Supplie, dict_to_typed
+from typing import List, Literal
+from datadis.types import (
+    ConsumptionData, ContractDetail, Supplie, dict_to_typed)
 import requests
 
 _HOST = 'https://datadis.es'
@@ -7,6 +8,7 @@
     'get_token': f'{_HOST}/nikola-auth/tokens/login',
     'get_supplies': f'{_HOST}/api-private/api/get-supplies',
     'get_contract_detail': f'{_HOST}/api-private/api/get-contract-detail',
+    'get_consumption_data': f'{_HOST}/api-private/api/get-consumption-data',
 }
 
 
@@ -82,3 +84,23 @@ def get_contract_detail(token: str, cups: str,
         return result
     else:
         raise ConnectionError(f'Error: {r.json()["message"]}')
+
+
+def get_consumption_data(token: str, cups: str,
+                         distrubutor_code: str, start_date: str, end_date: str,
+                         measurement_type: Literal[0, 1]) -> List[ConsumptionData]:
+    headers = {'Authorization': f'Bearer {token}'}
+
+    r = requests.get(_ENDPOINTS['get_consumption_data']
+                     + f'?cups={cups}&distributorCode={distrubutor_code}'
+                     + f'&start_date={start_date}&end_date={end_date}'
+                     + f'&measurement_type={measurement_type}',
+                     headers=headers)
+
+    if r.status_code == 200:
+        result = []
+        for contract in r.json():
+            result.append(dict_to_typed(contract, ConsumptionData))
+        return result
+    else:
+        raise ConnectionError(f'Error: {r.json()["message"]}')
diff --git a/datadis/types.py b/datadis/types.py
index fec007c..c549ee9 100644
--- a/datadis/types.py
+++ b/datadis/types.py
@@ -30,9 +30,17 @@ class ContractDetail(TypedDict):
     endDate: str
 
 
+class ConsumptionData(TypedDict):
+    cups: str
+    date: str
+    time: str
+    consumptionKWh: float
+    obtainMethod: str
+
+
 def dict_to_typed(data: Mapping[str, Any], typed: TypedDict) -> TypedDict:
     result = typed()
-    for key, key_type in typed.__annotations__.items():
+    for key, _ in typed.__annotations__.items():
         if key not in data:
             raise ValueError(f"Key: {key} is not available in data.")
         result[key] = data[key]
diff --git a/tests/test_datadis.py b/tests/test_datadis.py
index 3f81aa6..d8d40cf 100644
--- a/tests/test_datadis.py
+++ b/tests/test_datadis.py
@@ -1,4 +1,5 @@
-from datadis import get_token, get_supplies, get_contract_detail, _ENDPOINTS
+from datadis import (get_consumption_data, get_token,
+                     get_supplies, get_contract_detail, _ENDPOINTS)
 from unittest import mock
 
 
@@ -32,7 +33,7 @@ def text(self):
     elif args[0].startswith(_ENDPOINTS['get_contract_detail']):
         return MockResponse([{
             "address": "home",
-            "cups": "1234ABC",
+            "cups": "c",
             "postalCode": "1024",
             "province": "madrid",
             "municipality": "madrid",
@@ -49,6 +50,14 @@ def text(self):
             "startDate": "2020/09",
             "endDate": "2020/09",
         }], status_code=200)
+    elif args[0].startswith(_ENDPOINTS['get_consumption_data']):
+        return MockResponse([{
+            "cups": "1234ABC",
+            "date": "2021/08/01",
+            "time": "01:00",
+            "consumptionKWh": 0.194,
+            "obtainMethod": "Real"
+        }])
     else:
         return MockResponse([{"cups": "of rice"}], 'token_ok', 200)
 
@@ -73,3 +82,11 @@ def test_get_contract_detail(mock_get: mock.MagicMock):
     contract_detail = get_contract_detail("token", "cupaso", 2)
 
     assert contract_detail is not None
+
+
+@mock.patch('requests.get', side_effect=mock_requests)
+def test_get_consumption_data(mock_get: mock.MagicMock):
+    contract_detail = get_consumption_data(
+        "token", "cupaso", '2', '2021/08/01', '2021/08/31', 0)
+
+    assert contract_detail is not None