Skip to content

Commit

Permalink
fix(api-client): Updated get_account to use graphql so that inactive …
Browse files Browse the repository at this point in the history 
…meters are ignored
  • Loading branch information
@BottlecapDave
BottlecapDave committed Jul 2, 2022
1 parent 4a85bae commit f05dcd9
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 21 deletions.
118 changes: 102 additions & 16 deletions custom_components/octopus_energy/api_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,64 @@

_LOGGER = logging.getLogger(__name__)

api_token_query = '''mutation {{
obtainKrakenToken(input: {{ APIKey: "{api_key}" }}) {{
token
}}
}}'''

account_query = '''query {{
account(accountNumber: "{account_id}") {{
electricityAgreements(active: true) {{
meterPoint {{
mpan
meters(includeInactive: false) {{
serialNumber
smartExportElectricityMeter {{
deviceId
}}
}}
agreements {{
validFrom
validTo
tariff {{
...on StandardTariff {{
tariffCode
}}
...on DayNightTariff {{
tariffCode
}}
...on ThreeRateTariff {{
tariffCode
}}
...on HalfHourlyTariff {{
tariffCode
}}
...on PrepayTariff {{
tariffCode
}}
}}
}}
}}
}}
gasAgreements(active: true) {{
meterPoint {{
mprn
meters(includeInactive: false) {{
serialNumber
}}
agreements {{
validFrom
validTo
tariff {{
tariffCode
}}
}}
}}
}}
}}
}}'''

class OctopusEnergyApiClient:

def __init__(self, api_key):
Expand All @@ -22,23 +80,51 @@ def __init__(self, api_key):
async def async_get_account(self, account_id):
"""Get the user's account"""
async with aiohttp.ClientSession() as client:
auth = aiohttp.BasicAuth(self._api_key, '')
url = f'{self._base_url}/v1/accounts/{account_id}'
async with client.get(url, auth=auth) as response:
data = await self.__async_read_response(response, url, None)
if (data != None and "properties" in data):
# We're only supporting one property at the moment and we don't want to expose addresses
properties = data["properties"]
prop = next(current_prop for current_prop in properties if current_prop["moved_out_at"] == None)
if (prop == None):
raise Exception("Failed to find occupied property")

return {
"electricity_meter_points": prop["electricity_meter_points"],
"gas_meter_points": prop["gas_meter_points"]
}
url = f'{self._base_url}/v1/graphql/'
payload = { "query": api_token_query.format(api_key=self._api_key) }
async with client.post(url, json=payload) as token_response:
token_response_body = await self.__async_read_response(token_response, url, None)
if (token_response_body != None and "data" in token_response_body):
token = token_response_body["data"]["obtainKrakenToken"]["token"]

# Get account response
payload = { "query": account_query.format(account_id=account_id) }
headers = { "Authorization": f"JWT {token}" }
async with client.post(url, json=payload, headers=headers) as account_response:
account_response_body = await self.__async_read_response(account_response, url, None)
if (account_response_body != None and "data" in account_response_body):
return {
"electricity_meter_points": list(map(lambda mp: {
"mpan": mp["meterPoint"]["mpan"],
"meters": list(map(lambda m: {
"serial_number": m["serialNumber"],
"is_export": m["smartExportElectricityMeter"] != None,
}, mp["meterPoint"]["meters"])),
"agreements": list(map(lambda a: {
"valid_from": a["validFrom"],
"valid_to": a["validTo"],
"tariff_code": a["tariff"]["tariffCode"],
}, mp["meterPoint"]["agreements"]))
}, account_response_body["data"]["account"]["electricityAgreements"])),
"gas_meter_points": list(map(lambda mp: {
"mprn": mp["meterPoint"]["mprn"],
"meters": list(map(lambda m: {
"serial_number": m["serialNumber"],
}, mp["meterPoint"]["meters"])),
"agreements": list(map(lambda a: {
"valid_from": a["validFrom"],
"valid_to": a["validTo"],
"tariff_code": a["tariff"]["tariffCode"],
}, mp["meterPoint"]["agreements"]))
}, account_response_body["data"]["account"]["gasAgreements"])),
}
else:
_LOGGER.error("Failed to retrieve account")

return None
else:
_LOGGER.error("Failed to retrieve auth token")

return None

async def async_get_electricity_standard_rates(self, product_code, tariff_code, period_from, period_to):
"""Get the current standard rates"""
Expand Down
8 changes: 4 additions & 4 deletions custom_components/octopus_energy/sensor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -126,10 +126,10 @@ async def async_setup_default_sensors(hass, entry, async_add_entities):
for meter in point["meters"]:
_LOGGER.info(f'Adding electricity meter; mpan: {point["mpan"]}; serial number: {meter["serial_number"]}')
coordinator = create_reading_coordinator(hass, client, True, point["mpan"], meter["serial_number"])
entities.append(OctopusEnergyPreviousAccumulativeElectricityReading(coordinator, point["mpan"], meter["serial_number"], point["is_export"]))
entities.append(OctopusEnergyPreviousAccumulativeElectricityCost(coordinator, client, electricity_tariff_code, point["mpan"], meter["serial_number"], point["is_export"]))
entities.append(OctopusEnergyElectricityCurrentRate(rate_coordinator, point["mpan"], meter["serial_number"], point["is_export"]))
entities.append(OctopusEnergyElectricityPreviousRate(rate_coordinator, point["mpan"], meter["serial_number"], point["is_export"]))
entities.append(OctopusEnergyPreviousAccumulativeElectricityReading(coordinator, point["mpan"], meter["serial_number"], meter["is_export"]))
entities.append(OctopusEnergyPreviousAccumulativeElectricityCost(coordinator, client, electricity_tariff_code, point["mpan"], meter["serial_number"], meter["is_export"]))
entities.append(OctopusEnergyElectricityCurrentRate(rate_coordinator, point["mpan"], meter["serial_number"], meter["is_export"]))
entities.append(OctopusEnergyElectricityPreviousRate(rate_coordinator, point["mpan"], meter["serial_number"], meter["is_export"]))
else:
for meter in point["meters"]:
_LOGGER.info(f'Skipping electricity meter due to no active agreement; mpan: {point["mpan"]}; serial number: {meter["serial_number"]}')
Expand Down
3 changes: 2 additions & 1 deletion tests/integration/api_client/test_get_account.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,16 @@ async def test_when_get_account_is_called_then_electricity_and_gas_points_return
account = await client.async_get_account(account_id)

# Assert
assert account != None
assert "electricity_meter_points" in account

assert len(account["electricity_meter_points"]) == 1
meter_point = account["electricity_meter_points"][0]
assert meter_point["mpan"] == context["electricity_mpan"]
assert meter_point["is_export"] == False

assert len(meter_point["meters"]) == 1
meter = meter_point["meters"][0]
assert meter["is_export"] == False
assert meter["serial_number"] == context["electricity_serial_number"]

assert "agreements" in meter_point
Expand Down

0 comments on commit f05dcd9

Please sign in to comment.