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