def test_json_output(capsys):
"""Test json output function."""
data = {
"overview": {
"contract_id": "foo_customer",
"account_id": "foo_accoutn",
"customer_id": "foo_id",
"balance": "foo_balance"
},
"current_period": {
"period_total_bill": "foobar_total_bill", "current_period_data": "current_period_data"
},
"current_annual_data": {
"some_annual_data_key": "foobar_annual_data_value"
},
"yesterday_data": {
"date": "some_date",
"hour1": {
"foobar": "some_data"
}
}
}
# \n because of trailing newline in readouterr
expected = f'{json.dumps(data)}\n'
class MockCustomer: # pylint: disable=too-few-public-methods
"""Mock class for Customer."""
contract_id = "foo_customer"
account_id = "foo_accoutn"
customer_id = "foo_id"
balance = "foo_balance"
current_period = {
"period_total_bill": "foobar_total_bill", "current_period_data": "current_period_data"
}
current_annual_data = {"some_annual_data_key": "foobar_annual_data_value"}
current_daily_data = {"some_date": {"hour1": {"foobar": "some_data"}}}
mock_customer = MockCustomer()
output_json(mock_customer)
captured = capsys.readouterr()
assert captured.out == expected
def main():
"""Entrypoint function."""
parser = argparse.ArgumentParser()
parser.add_argument('-u', '--username', help='Hydro Quebec username')
parser.add_argument('-p', '--password', help='Password')
parser.add_argument('-j',
'--json',
action='store_true',
default=False,
help='Json output')
parser.add_argument('-i',
'--influxdb',
action='store_true',
default=False,
help='InfluxDb output')
parser.add_argument('-c',
'--contract',
default=None,
help='Contract number')
parser.add_argument('--client-number',
default=None,
help='Use client Number to speed up login')
parser.add_argument('-l',
'--list-contracts',
action='store_true',
default=False,
help='List all your contracts')
parser.add_argument('-H',
'--hourly',
action='store_true',
default=False,
help='Show yesterday hourly consumption')
parser.add_argument('-D',
'--dump-data',
action='store_true',
default=False,
help='Show contract python object as dict')
parser.add_argument('-t',
'--timeout',
default=REQUESTS_TIMEOUT,
help='Request timeout')
parser.add_argument(
'-L',
'--log-level',
choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
default='WARNING',
help='Log level')
parser.add_argument('-V',
'--version',
action='store_true',
default=False,
help='Show version')
raw_group = parser.add_argument_group(
'Detailled-energy raw download option')
raw_group.add_argument('--detailled-energy',
action='store_true',
default=False,
help='Get raw json output download')
raw_group.add_argument('--start-date',
default=(datetime.now(HQ_TIMEZONE) -
timedelta(days=1)).strftime("%Y-%m-%d"),
help='Start date for detailled-output')
raw_group.add_argument(
'--end-date',
default=datetime.now(HQ_TIMEZONE).strftime("%Y-%m-%d"),
help="End date for detailled-output")
args = parser.parse_args()
if args.version:
print(VERSION)
return 0
# Check input for Username, Password and Contract - CLI overwrite ENV variable
# Check Env
hydro_user = os.environ.get("PYHQ_USER")
hydro_pass = os.environ.get("PYHQ_PASSWORD")
hydro_contract = os.environ.get("PYHQ_CONTRACT")
hydro_client_number = os.environ.get("PYHQ_CLIENT_NUMBER")
# Check Cli
if args.username:
hydro_user = args.username
if args.password:
hydro_pass = args.password
if args.contract:
hydro_contract = args.contract
if args.client_number:
hydro_client_number = args.client_number
if not hydro_user or not hydro_pass:
parser.print_usage()
print("pyhydroquebec: error: the following arguments are required: "
"-u/--username, -p/--password")
return 3
client = HydroQuebecClient(hydro_user,
hydro_pass,
args.timeout,
log_level=args.log_level)
loop = asyncio.get_event_loop()
# Get the async_func
if args.list_contracts:
async_func = list_contracts(client)
elif args.dump_data:
async_func = dump_data(client, hydro_contract)
elif args.detailled_energy is False:
async_func = fetch_data(client, hydro_contract, hydro_client_number,
args.hourly)
else:
start_date = datetime.strptime(args.start_date, '%Y-%m-%d')
end_date = datetime.strptime(args.end_date, '%Y-%m-%d')
async_func = fetch_data_detailled_energy_use(client, hydro_contract,
hydro_client_number,
start_date, end_date)
# Fetch data
try:
results = loop.run_until_complete(asyncio.gather(async_func))
except BaseException as exp:
print(exp)
print(traceback.print_tb(exp.__traceback__))
return 1
finally:
close_fut = asyncio.wait([client.close_session()])
loop.run_until_complete(close_fut)
loop.close()
# Output data
if args.list_contracts:
for customer in results[0]:
print("Contract: {contract_id}\n\t"
"Account: {account_id}\n\t"
"Customer: {customer_id}".format(**customer))
elif args.dump_data or args.detailled_energy:
pprint(results[0].__dict__)
elif args.influxdb:
output_influx(results[0])
elif args.json:
output_json(results[0])
else:
output_text(results[0], args.hourly)
return 0