# # 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. # # Documentation for orm: # https://github.com/encode/orm # # To interactively work with a database: # import bitelab.data # import databases # database = databases.Database('sqlite:///' + dbpath) # data = bitelab.data.make_orm(database) from typing import Optional, Union, Dict, Any from pydantic import BaseModel, Field from datetime import datetime import databases import orm import sqlalchemy import tempfile import unittest __all__ = [ 'make_orm', 'BoardClassInfo', 'Board', 'Error' ] class BoardClassInfo(BaseModel): clsname: str arch: str class Board(BaseModel): name: str brdclass: str reserved: bool attrs: Dict[str, Any] = Field(default_factory=dict) class Config: orm_mode = True class Error(BaseModel): error: str board: Optional[Board] def _issubclass(a, b): try: return issubclass(a, b) except TypeError: return False class DataWrapper(object): pass def make_orm(database): metadata = sqlalchemy.MetaData() class BoardStatus(orm.Model): __tablename__ = 'boardstatus' __database__ = database __metadata__ = metadata board = orm.Text(primary_key=True) user = orm.Text(index=True) time_reserved = orm.DateTime(default=lambda: datetime.utcnow()) class APIKey(orm.Model): __tablename__ = 'apikeys' __database__ = database __metadata__ = metadata user = orm.Text(index=True) key = orm.Text(primary_key=True) engine = sqlalchemy.create_engine(str(database.url)) metadata.create_all(engine) r = DataWrapper() lcls = locals() for i in [ 'engine', 'metadata', ] + [ x for x in lcls if _issubclass(lcls[x], orm.Model) ]: setattr(r, i, lcls[i]) return r async def _setup_data(data): await data.APIKey.objects.create(user='foo', key='thisisanapikey') await data.APIKey.objects.create(user='bar', key='anotherlongapikey') class TestDatabase(unittest.IsolatedAsyncioTestCase): def setUp(self): # setup temporary directory self.dbtempfile = tempfile.NamedTemporaryFile() self.database = databases.Database('sqlite:///' + self.dbtempfile.name) self.data = make_orm(self.database) def tearDown(self): self.data = None self.database = None self.dbtempfile = None async def test_apikey(self): data = self.data # that the test database starts empty self.assertEqual(await data.APIKey.objects.all(), []) # that when it is populated with test data await _setup_data(data) # the data can be accessed self.assertEqual((await data.APIKey.objects.get( key='thisisanapikey')).user, 'foo') self.assertEqual((await data.APIKey.objects.get( key='anotherlongapikey')).user, 'bar')