move the main client to __main__..

it'll be getting more complex soon, and I'm trying to simplify __init__

bitelab/__init__.py

@@ -34,12 +34,12 @@ from io import StringIO
 from fastapi import APIRouter, Depends, FastAPI, HTTPException, Request
 from fastapi.security import OAuth2PasswordBearer
 from httpx import AsyncClient, Auth
-from unittest.mock import patch, AsyncMock, Mock, PropertyMock
 from starlette.responses import JSONResponse
 from starlette.status import HTTP_200_OK
 from starlette.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, \
     HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND, HTTP_409_CONFLICT
 from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
+from unittest.mock import patch, AsyncMock, Mock, PropertyMock
 
 from . import config
 from .data import *
@@ -336,6 +336,10 @@ async def release_board(board_id, user: str = Depends(lookup_user),
 
 	return brd
 
+@router.post('/board/{board_id_or_class}/attrs', response_model=Union[Board, Error])
+async def set_board_attrs():
+	pass
+
 @router.get('/board/',response_model=Dict[str, Board])
 async def get_boards(user: str = Depends(lookup_user),
     brdmgr: BoardManager = Depends(get_boardmanager)):
@@ -363,51 +367,6 @@ def getApp():
 # uvicorn can't call the above function, while hypercorn can
 #app = getApp()
 
-def check_res_code(res):
-	if res.status_code == HTTP_401_UNAUTHORIZED:
-		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()))
-		sys.exit(1)
-
-async def real_main():
-	baseurl = os.environ['BITELAB_URL']
-	authkey = os.environ['BITELAB_AUTH']
-
-	client = AsyncClient(base_url=baseurl)
-
-	try:
-		if sys.argv[1] == 'list':
-			res = await client.get('board/classes', auth=BiteAuth(authkey))
-
-			check_res_code(res)
-
-			print('Classes:')
-			for i in res.json():
-				print('\t' + i)
-
-			res.close()
-		elif sys.argv[1] in ('reserve', 'release'):
-			res = await client.post('board/%s/%s' %
-			    (urllib.parse.quote(sys.argv[2], safe=''),
-			     sys.argv[1]),
-			    auth=BiteAuth(authkey))
-
-			check_res_code(res)
-
-			brd = Board.parse_obj(res.json())
-			print('Name:\t%s' % brd.name)
-			print('Class:\t%s' % brd.brdclass)
-			print('Attributes:')
-			for i in sorted(brd.attrs):
-				print('\t%s\t%s' % (i, brd.attrs[i]))
-	finally:
-		await client.aclose()
-
-def main():
-	asyncio.run(real_main())
-
 class TestUnhashLRU(unittest.TestCase):
 	def test_unhashlru(self):
 		lsta = []
@@ -705,131 +664,3 @@ class TestDatabase(unittest.IsolatedAsyncioTestCase):
 		    key='thisisanapikey')).user, 'foo')
 		self.assertEqual((await data.APIKey.objects.get(
 		    key='anotherlongapikey')).user, 'bar')
-
-@patch.dict(os.environ, dict(BITELAB_URL='http://someserver/'))
-@patch.dict(os.environ, dict(BITELAB_AUTH='thisisanapikey'))
-class TestClient(unittest.TestCase):
-	def setUp(self):
-		self.ac_patcher = patch(__name__ + '.AsyncClient')
-		self.ac = self.ac_patcher.start()
-		self.addCleanup(self.ac_patcher.stop)
-
-		self.acg = self.ac.return_value.get = AsyncMock()
-		self.acaclose = self.ac.return_value.aclose = AsyncMock()
-		self.acgr = self.acg.return_value = Mock()
-
-		self.acp = self.ac.return_value.post = AsyncMock()
-		self.acpr = self.acp.return_value = Mock()
-
-	def runMain(self):
-		try:
-			stdout = StringIO()
-			with patch.dict(sys.__dict__, dict(stdout=stdout)):
-				main()
-
-			ret = 0
-		except SystemExit as e:
-			ret = e.code
-
-		return ret, stdout.getvalue()
-
-	@patch.dict(sys.__dict__, dict(argv=[ '', 'list' ]))
-	def test_list_failure(self):
-		ac = self.ac
-		acg = self.acg
-		acg.return_value.status_code = HTTP_401_UNAUTHORIZED
-		acg.return_value.json.return_value = {
-		    'detail': 'Invalid authentication credentials'
-		}
-
-		ret, stdout = self.runMain()
-
-		output = '''Invalid authentication credentials.
-'''
-
-		self.assertEqual(ret, 1)
-		# XXX -- really should go to stderr
-		self.assertEqual(stdout, output)
-
-		ac.assert_called_with(base_url='http://someserver/')
-
-		acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'))
-
-	# XXX - add error cases for UI
-
-	@patch.dict(sys.__dict__, dict(argv=[ '', 'list' ]))
-	def test_list(self):
-		ac = self.ac
-		acg = self.acg
-		acg.return_value.status_code = HTTP_200_OK
-		acg.return_value.json.return_value = { 'cora-z7s': {
-		    'arch': 'arm-armv7', 'clsname': 'cora-z7s', }}
-
-		ret, stdout = self.runMain()
-
-		output = '''Classes:
-	cora-z7s
-'''
-
-		self.assertEqual(ret, 0)
-		self.assertEqual(stdout, output)
-
-		ac.assert_called_with(base_url='http://someserver/')
-
-		acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'))
-
-	# XXX - add error cases for UI
-
-	@patch.dict(sys.__dict__, dict(argv=[ '', 'reserve', 'cora-z7s' ]))
-	def test_reserve(self):
-		ac = self.ac
-		acp = self.acp
-		acp.return_value.status_code = HTTP_200_OK
-		acp.return_value.json.return_value = Board(name='cora-1',
-		    brdclass='cora-z7s', reserved=True,
-		    attrs={
-		        'ip': '172.20.20.5',
-		        'power': False,
-		    }).dict()
-
-		ret, stdout = self.runMain()
-
-		output = '''Name:\tcora-1
-Class:\tcora-z7s
-Attributes:
-\tip\t172.20.20.5
-\tpower\tFalse
-'''
-
-		self.assertEqual(ret, 0)
-		self.assertEqual(stdout, output)
-
-		ac.assert_called_with(base_url='http://someserver/')
-
-		acp.assert_called_with('board/cora-z7s/reserve', auth=BiteAuth('thisisanapikey'))
-
-	@patch.dict(sys.__dict__, dict(argv=[ '', 'release', 'cora-z7s' ]))
-	def test_release(self):
-		ac = self.ac
-		acp = self.acp
-		acp.return_value.status_code = HTTP_200_OK
-		acp.return_value.json.return_value = Board(name='cora-1',
-		    brdclass='cora-z7s', reserved=False,
-		    attrs={
-		        'power': False,
-		    }).dict()
-
-		ret, stdout = self.runMain()
-
-		output = '''Name:\tcora-1
-Class:\tcora-z7s
-Attributes:
-\tpower\tFalse
-'''
-
-		self.assertEqual(ret, 0)
-		self.assertEqual(stdout, output)
-
-		ac.assert_called_with(base_url='http://someserver/')
-
-		acp.assert_called_with('board/cora-z7s/release', auth=BiteAuth('thisisanapikey'))

bitelab/__main__.py

@@ -1,4 +1,218 @@
-from . import main
+#
+# Copyright (c) 2020 The FreeBSD Foundation
+#
+# This software1 was developed by John-Mark Gurney under sponsorship
+# from the FreeBSD Foundation.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+# 1.  Redistributions of source code must retain the above copyright
+#     notice, this list of conditions and the following disclaimer.
+# 2.  Redistributions in binary form must reproduce the above copyright
+#     notice, this list of conditions and the following disclaimer in the
+#     documentation and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
+# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
+# SUCH DAMAGE.
+#
 
-if __name__ == '__main__':
+from httpx import AsyncClient, Auth
+from io import StringIO
+from starlette.status import HTTP_200_OK
+from starlette.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, \
+    HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND, HTTP_409_CONFLICT
+from unittest.mock import patch, AsyncMock, Mock
+
+from . import BiteAuth, Board
+
+import asyncio
+import os
+import sys
+import unittest
+import urllib
+
+def check_res_code(res):
+	if res.status_code == HTTP_401_UNAUTHORIZED:
+		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()))
+		sys.exit(1)
+
+async def real_main():
+	baseurl = os.environ['BITELAB_URL']
+	authkey = os.environ['BITELAB_AUTH']
+
+	client = AsyncClient(base_url=baseurl)
+
+	try:
+		if sys.argv[1] == 'list':
+			res = await client.get('board/classes', auth=BiteAuth(authkey))
+
+			check_res_code(res)
+
+			print('Classes:')
+			for i in res.json():
+				print('\t' + i)
+
+			res.close()
+		elif sys.argv[1] in ('reserve', 'release'):
+			res = await client.post('board/%s/%s' %
+			    (urllib.parse.quote(sys.argv[2], safe=''),
+			     sys.argv[1]),
+			    auth=BiteAuth(authkey))
+
+			check_res_code(res)
+
+			brd = Board.parse_obj(res.json())
+			print('Name:\t%s' % brd.name)
+			print('Class:\t%s' % brd.brdclass)
+			print('Attributes:')
+			for i in sorted(brd.attrs):
+				print('\t%s\t%s' % (i, brd.attrs[i]))
+	finally:
+		await client.aclose()
+
+def main():
+	asyncio.run(real_main())
+
+if __name__ == '__main__': #pragma: no cover
 	main()
+
+@patch.dict(os.environ, dict(BITELAB_URL='http://someserver/'))
+@patch.dict(os.environ, dict(BITELAB_AUTH='thisisanapikey'))
+class TestClient(unittest.TestCase):
+	def setUp(self):
+		self.ac_patcher = patch(__name__ + '.AsyncClient')
+		self.ac = self.ac_patcher.start()
+		self.addCleanup(self.ac_patcher.stop)
+
+		self.acg = self.ac.return_value.get = AsyncMock()
+		self.acaclose = self.ac.return_value.aclose = AsyncMock()
+		self.acgr = self.acg.return_value = Mock()
+
+		self.acp = self.ac.return_value.post = AsyncMock()
+		self.acpr = self.acp.return_value = Mock()
+
+	def runMain(self):
+		try:
+			stdout = StringIO()
+			with patch.dict(sys.__dict__, dict(stdout=stdout)):
+				main()
+
+			ret = 0
+		except SystemExit as e:
+			ret = e.code
+
+		return ret, stdout.getvalue()
+
+	@patch.dict(sys.__dict__, dict(argv=[ '', 'list' ]))
+	def test_list_failure(self):
+		ac = self.ac
+		acg = self.acg
+		acg.return_value.status_code = HTTP_401_UNAUTHORIZED
+		acg.return_value.json.return_value = {
+		    'detail': 'Invalid authentication credentials'
+		}
+
+		ret, stdout = self.runMain()
+
+		output = '''Invalid authentication credentials.
+'''
+
+		self.assertEqual(ret, 1)
+		# XXX -- really should go to stderr
+		self.assertEqual(stdout, output)
+
+		ac.assert_called_with(base_url='http://someserver/')
+
+		acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'))
+
+	# XXX - add error cases for UI
+
+	@patch.dict(sys.__dict__, dict(argv=[ '', 'list' ]))
+	def test_list(self):
+		ac = self.ac
+		acg = self.acg
+		acg.return_value.status_code = HTTP_200_OK
+		acg.return_value.json.return_value = { 'cora-z7s': {
+		    'arch': 'arm-armv7', 'clsname': 'cora-z7s', }}
+
+		ret, stdout = self.runMain()
+
+		output = '''Classes:
+	cora-z7s
+'''
+
+		self.assertEqual(ret, 0)
+		self.assertEqual(stdout, output)
+
+		ac.assert_called_with(base_url='http://someserver/')
+
+		acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'))
+
+	# XXX - add error cases for UI
+
+	@patch.dict(sys.__dict__, dict(argv=[ '', 'reserve', 'cora-z7s' ]))
+	def test_reserve(self):
+		ac = self.ac
+		acp = self.acp
+		acp.return_value.status_code = HTTP_200_OK
+		acp.return_value.json.return_value = Board(name='cora-1',
+		    brdclass='cora-z7s', reserved=True,
+		    attrs={
+		        'ip': '172.20.20.5',
+		        'power': False,
+		    }).dict()
+
+		ret, stdout = self.runMain()
+
+		output = '''Name:\tcora-1
+Class:\tcora-z7s
+Attributes:
+\tip\t172.20.20.5
+\tpower\tFalse
+'''
+
+		self.assertEqual(ret, 0)
+		self.assertEqual(stdout, output)
+
+		ac.assert_called_with(base_url='http://someserver/')
+
+		acp.assert_called_with('board/cora-z7s/reserve', auth=BiteAuth('thisisanapikey'))
+
+	@patch.dict(sys.__dict__, dict(argv=[ '', 'release', 'cora-z7s' ]))
+	def test_release(self):
+		ac = self.ac
+		acp = self.acp
+		acp.return_value.status_code = HTTP_200_OK
+		acp.return_value.json.return_value = Board(name='cora-1',
+		    brdclass='cora-z7s', reserved=False,
+		    attrs={
+		        'power': False,
+		    }).dict()
+
+		ret, stdout = self.runMain()
+
+		output = '''Name:\tcora-1
+Class:\tcora-z7s
+Attributes:
+\tpower\tFalse
+'''
+
+		self.assertEqual(ret, 0)
+		self.assertEqual(stdout, output)
+
+		ac.assert_called_with(base_url='http://someserver/')
+
+		acp.assert_called_with('board/cora-z7s/release', auth=BiteAuth('thisisanapikey'))

bitelab/testing.py

@@ -29,4 +29,5 @@
 # Module that includes all the test cases.
 
 from .snmp import TestSNMPPower, TestSNMPWrapper
-from . import TestClient, TestDatabase, TestBiteLab, TestUnhashLRU
+from . import TestDatabase, TestBiteLab, TestUnhashLRU
+from .__main__ import TestClient