From 6a1b85215d635c375014a21874cc186e815c0e89 Mon Sep 17 00:00:00 2001 From: John-Mark Gurney Date: Tue, 8 Dec 2020 23:58:52 -0800 Subject: [PATCH] move arg parsing a bit earlier, handle non-json error responses... --- bitelab/__main__.py | 32 ++++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/bitelab/__main__.py b/bitelab/__main__.py index 95a2343..e1940e9 100644 --- a/bitelab/__main__.py +++ b/bitelab/__main__.py @@ -39,6 +39,7 @@ import argparse import asyncio import contextlib import io +import json import os import sys import unittest @@ -49,7 +50,11 @@ def check_res_code(res): print('Invalid authentication credentials.') sys.exit(1) elif res.status_code != HTTP_200_OK: - print('Got status: %d, json: %s' % (res.status_code, res.json())) + try: + print('Got status: %d, json: %s' % (res.status_code, res.json())) + except json.decoder.JSONDecodeError: + # body is JSON + print('Got status: %d, body: %s' % (res.status_code, repr(res.text))) sys.exit(1) def makebool(s): @@ -94,11 +99,6 @@ def get_sshpubkey(fname): raise IOError async def real_main(): - baseurl = os.environ['BITELAB_URL'] - authkey = os.environ['BITELAB_AUTH'] - - client = AsyncClient(base_url=baseurl) - parser = argparse.ArgumentParser() subparsers = parser.add_subparsers(title='subcommands', dest='subparser_name', @@ -117,6 +117,11 @@ async def real_main(): args = parser.parse_args() #print(repr(args), file=sys.stderr) + baseurl = os.environ['BITELAB_URL'] + authkey = os.environ['BITELAB_AUTH'] + + client = AsyncClient(base_url=baseurl) + try: if args.subparser_name == 'list': res = await client.get('board/classes', @@ -179,11 +184,11 @@ class TestClient(unittest.TestCase): self.acp = self.ac.return_value.post = AsyncMock() self.acpr = self.acp.return_value = Mock() - def runMain(self): + def runMain(self, fun=main): try: stdout = StringIO() with patch.dict(sys.__dict__, dict(stdout=stdout)): - main() + fun() ret = 0 except SystemExit as e: @@ -393,5 +398,12 @@ Attributes: self.assertRaises(ValueError, make_attrs, 'power=bar') def test_check_res_code(self): - # XXX - on weird errors, res.json() will fail w/ json.decoder.JSONDecodeError - pass + res = Mock() + res.status_code = HTTP_404_NOT_FOUND + res.json.side_effect = json.decoder.JSONDecodeError('foo', 'bar', 1) + res.text = 'body' + + ret, output = self.runMain(lambda: check_res_code(res)) + + self.assertEqual(ret, 1) + self.assertEqual(output, 'Got status: 404, body: \'body\'\n')