diff --git a/bitelab/__main__.py b/bitelab/__main__.py index b9c6a50..f3bc49a 100644 --- a/bitelab/__main__.py +++ b/bitelab/__main__.py @@ -55,10 +55,12 @@ def check_res_code(res): sys.exit(1) elif res.status_code != HTTP_200_OK: try: - print('Got status: %d, json: %s' % (res.status_code, res.json())) + 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))) + print('Got status: %d, body: %s' % (res.status_code, + repr(res.text))) sys.exit(1) def makebool(s): @@ -125,18 +127,22 @@ async def fwd_data(reader, writer): await writer.drain() async def run_exec(baseurl, authkey, board, args): - url = urllib.parse.urljoin(baseurl, 'board/%s/exec' % urllib.parse.quote(board, safe='')) + url = urllib.parse.urljoin(baseurl, 'board/%s/exec' % + urllib.parse.quote(board, safe='')) url = convert_to_ws(url) stdin, stdout = await aioconsole.stream.get_standard_streams() - async with websockets.connect(url) as ws, wsfwd.WSFWDClient(ws.recv, ws.send) as client: + async with websockets.connect(url) as ws, wsfwd.WSFWDClient(ws.recv, + ws.send) as client: try: await client.auth(dict(bearer=authkey)) proc = await client.exec(args=args) - toexec_task = asyncio.create_task(fwd_data(stdin, proc.stdin)) - fromexec_task = asyncio.create_task(fwd_data(proc.stdout, stdout)) + toexec_task = asyncio.create_task(fwd_data(stdin, + proc.stdin)) + fromexec_task = asyncio.create_task(fwd_data( + proc.stdout, stdout)) r = await proc.wait() @@ -159,20 +165,30 @@ async def real_main(): dest='subparser_name', description='valid subcommands', help='additional help') - parse_list = subparsers.add_parser('list', help='list available board classes') - - parser_reserve = subparsers.add_parser('reserve', aliases=[ 'release' ], help='reserve/release a board') - parser_reserve.add_argument('-i', metavar='identity_file', type=str, help='file name for ssh public key') - parser_reserve.add_argument('board', type=str, help='name of the board or class') - - parser_set = subparsers.add_parser('set', help='set attributes on a board') - parser_set.add_argument('setvars', type=str, nargs='+', help='name of the board or class') - parser_set.add_argument('board', type=str, help='name of the board or class') - - parser_exec = subparsers.add_parser('exec', help='run a program in the jail for a board') - parser_exec.add_argument('board', type=str, help='name of the board or class') + parse_list = subparsers.add_parser('list', + help='list available board classes') + + parser_reserve = subparsers.add_parser('reserve', aliases=[ 'release' ], + help='reserve/release a board') + parser_reserve.add_argument('-i', metavar='identity_file', type=str, + help='file name for ssh public key') + parser_reserve.add_argument('board', type=str, + help='name of the board or class') + + parser_set = subparsers.add_parser('set', + help='set attributes on a board') + parser_set.add_argument('setvars', type=str, nargs='+', + help='name of the board or class') + parser_set.add_argument('board', type=str, + help='name of the board or class') + + parser_exec = subparsers.add_parser('exec', + help='run a program in the jail for a board') + parser_exec.add_argument('board', type=str, + help='name of the board or class') parser_exec.add_argument('prog', type=str, help='program to exec') - parser_exec.add_argument('args', type=str, nargs='*', help='arguments for program') + parser_exec.add_argument('args', type=str, nargs='*', + help='arguments for program') args = parser.parse_args() #print(repr(args), file=sys.__stderr__) @@ -202,7 +218,8 @@ async def real_main(): elif args.subparser_name in ('reserve', 'release'): kwargs = _httpxargs.copy() with contextlib.suppress(IOError): - kwargs['json'] = dict(sshpubkey=get_sshpubkey(args.i)) + kwargs['json'] = dict(sshpubkey=get_sshpubkey( + args.i)) res = await client.post('board/%s/%s' % (urllib.parse.quote(args.board, safe=''), args.subparser_name), @@ -260,7 +277,8 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase): @contextlib.contextmanager def make_pipe(self): r, w = os.pipe() - with os.fdopen(r, 'rb', buffering=65536) as readfl, os.fdopen(w, 'wb', buffering=65536) as writefl: + with os.fdopen(r, 'rb', buffering=65536) as readfl, \ + os.fdopen(w, 'wb', buffering=65536) as writefl: yield readfl, writefl # too lazy to make this async since async file-like objects @@ -284,23 +302,30 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase): stdout = io.BytesIO() # Data path: - # stdin -> stdin_task -> stdinwriter -> pipe -> stdinreader -> sys.stdin -> real_main -> - # sys.stdout -> stdoutwriter -> pipe -> stdoutreader -> stdoud_task -> stdout + # stdin -> stdin_task -> stdinwriter -> pipe -> + # stdinreader -> sys.stdin -> real_main -> sys.stdout -> + # stdoutwriter -> pipe -> stdoutreader -> stdoud_task -> stdout # # How things get closed: - # stdin is already "closed", stdin_task will close stdinwriter when eof encountered (doclose). + # stdin is already "closed", stdin_task will close + # stdinwriter when eof encountered (doclose). # create the pipes needed - with self.make_pipe() as (stdinreader, stdinwriter), self.make_pipe() as (stdoutreader, stdoutwriter): + with self.make_pipe() as (stdinreader, stdinwriter), \ + self.make_pipe() as (stdoutreader, stdoutwriter): # setup the threads to move data loop = asyncio.get_running_loop() - stdin_task = loop.run_in_executor(None, self.copytask, io.BytesIO(stdin), stdinwriter) + stdin_task = loop.run_in_executor(None, self.copytask, + io.BytesIO(stdin), stdinwriter) - # do not close stdout, otherwise we cannot obtain the value - stdout_task = loop.run_in_executor(None, self.copytask, stdoutreader, stdout, False) + # do not close stdout, otherwise we cannot obtain + # the value + stdout_task = loop.run_in_executor(None, self.copytask, + stdoutreader, stdout, False) # insert the pipes - with patch.dict(sys.__dict__, dict(stdin=TextIOWrapper(stdinreader), + with patch.dict(sys.__dict__, + dict(stdin=TextIOWrapper(stdinreader), stdout=TextIOWrapper(stdoutwriter))): try: # run the function @@ -341,18 +366,21 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase): raise RuntimeError('badauth') server = TestServer(self.toserver.get, self.toclient.put) - with patch.dict(sys.__dict__, dict(argv=[ 'rand', 'exec', 'cora-z7s', 'program', 'arg1', 'arg2' ])), \ + with patch.dict(sys.__dict__, dict(argv=[ 'rand', 'exec', + 'cora-z7s', 'program', 'arg1', 'arg2' ])), \ patch('websockets.connect') as webcon: self.setup_websockets_mock(webcon) ret, stdout = await self.runAsyncMain() - webcon.assert_called_with(urllib.parse.urljoin('ws://someserver/', 'board/cora-z7s/exec')) + webcon.assert_called_with(urllib.parse.urljoin( + 'ws://someserver/', 'board/cora-z7s/exec')) await server.__aexit__(None, None, None) await asyncio.sleep(.1) - self.assertEqual(stdout.decode(), 'failed to exec: Got auth error: \'badauth\'\n') + self.assertEqual(stdout.decode(), + 'failed to exec: Got auth error: \'badauth\'\n') self.assertEqual(ret, 1) @@ -368,11 +396,13 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase): async def handle_chanclose(self, msg): self.add_tasks(asyncio.create_task(self.sendcmd( - dict(cmd='chanclose', chan=self._stdout_stream)))) + dict(cmd='chanclose', + chan=self._stdout_stream)))) async def handle_exec(self, msg): self._stdout_stream = msg['stdout'] - assert msg['args'] == [ 'program', 'arg1', 'arg2' ] + assert msg['args'] == [ 'program', 'arg1', + 'arg2' ] self.add_stream_handler(msg['stdin'], functools.partial(self.echo_handler, msg['stdout'])) @@ -383,7 +413,8 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase): server = TestServer(self.toserver.get, self.toclient.put) - with patch.dict(sys.__dict__, dict(argv=[ 'rand', 'exec', 'cora-z7s', 'program', 'arg1', 'arg2' ])), \ + with patch.dict(sys.__dict__, dict(argv=[ 'rand', 'exec', + 'cora-z7s', 'program', 'arg1', 'arg2' ])), \ patch('websockets.connect') as webcon: self.setup_websockets_mock(webcon) @@ -435,21 +466,25 @@ class TestClient(unittest.TestCase): ] for orig, new in testpairs: - self.assertEqual(convert_to_ws(orig), new, 'failed to convert: %s' % repr(orig)) + self.assertEqual(convert_to_ws(orig), new, + 'failed to convert: %s' % repr(orig)) def test_sshpubkey(self): fname = 'fname' rdata = 'foo' - with patch('builtins.open', mock_open(read_data=rdata)) as mock_file: + with patch('builtins.open', mock_open(read_data=rdata)) as \ + mock_file: self.assertEqual(get_sshpubkey(fname), rdata) mock_file.assert_called_with(fname) with patch('builtins.open') as mock_file: - mock_file.side_effect = [ IOError(), mock_open(read_data=rdata)() ] + mock_file.side_effect = [ IOError(), + mock_open(read_data=rdata)() ] self.assertEqual(get_sshpubkey(None), rdata) for i in ('id_ed25519', 'id_rsa'): - fname = os.path.expanduser(os.path.join('~', '.ssh', i + '.pub')) + fname = os.path.expanduser(os.path.join('~', + '.ssh', i + '.pub')) mock_file.assert_any_call(fname) with patch('builtins.open') as mock_file: @@ -477,7 +512,8 @@ class TestClient(unittest.TestCase): ac.assert_called_with(base_url='http://someserver/') - acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'), **_httpxargs) + acg.assert_called_with('board/classes', + auth=BiteAuth('thisisanapikey'), **_httpxargs) # XXX - add error cases for UI @@ -500,7 +536,8 @@ class TestClient(unittest.TestCase): ac.assert_called_with(base_url='http://someserver/') - acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'), **_httpxargs) + acg.assert_called_with('board/classes', + auth=BiteAuth('thisisanapikey'), **_httpxargs) # XXX - add error cases for UI @@ -542,7 +579,8 @@ Attributes: auth=BiteAuth('thisisanapikey'), **_httpxargs) @patch('bitelab.__main__.get_sshpubkey') - @patch.dict(sys.__dict__, dict(argv=[ '', 'reserve', '-i', 'bogusfilename', 'cora-z7s' ])) + @patch.dict(sys.__dict__, dict(argv=[ '', 'reserve', '-i', + 'bogusfilename', 'cora-z7s' ])) def test_reserve_ssh(self, gspk): ac = self.ac acp = self.acp @@ -606,8 +644,10 @@ Attributes: acp.assert_called_with('board/cora-z7s/release', auth=BiteAuth('thisisanapikey'), **_httpxargs) - @patch('bitelab.__main__._typeconv', dict(power=makebool, other=makebool)) - @patch.dict(sys.__dict__, dict(argv=[ '', 'set', 'power=on', 'other=off', 'cora-z7s' ])) + @patch('bitelab.__main__._typeconv', dict(power=makebool, + other=makebool)) + @patch.dict(sys.__dict__, dict(argv=[ '', 'set', 'power=on', + 'other=off', 'cora-z7s' ])) def test_set(self): ac = self.ac acp = self.acp @@ -634,8 +674,8 @@ Attributes: ac.assert_called_with(base_url='http://someserver/') acp.assert_called_with('board/cora-z7s/attrs', - auth=BiteAuth('thisisanapikey'), json=dict(power=True, other=False), - **_httpxargs) + auth=BiteAuth('thisisanapikey'), json=dict(power=True, + other=False), **_httpxargs) def test_make_attrs(self): self.assertEqual(make_attrs('power=on'), dict(power=True)) @@ -650,7 +690,8 @@ Attributes: def test_check_res_code(self): res = Mock() res.status_code = HTTP_404_NOT_FOUND - res.json.side_effect = json.decoder.JSONDecodeError('foo', 'bar', 1) + res.json.side_effect = json.decoder.JSONDecodeError('foo', + 'bar', 1) res.text = 'body' ret, output = self.runMain(lambda: check_res_code(res))