def test_add_column_with_index(self):
"""
Test adding a column with an index to a MigrationManager.
"""
manager = MigrationManager()
manager.add_column(
table_class_name="Manager",
tablename="manager",
column_name="email",
column_class=Varchar,
column_class_name="Varchar",
params={
"length": 100,
"default": "",
"null": True,
"primary": False,
"key": False,
"unique": True,
"index": True,
},
)
index_name = Manager._get_index_name(["email"])
asyncio.run(manager.run())
self.assertTrue(index_name in Manager.indexes().run_sync())
# Reverse
asyncio.run(manager.run_backwards())
self.assertTrue(index_name not in Manager.indexes().run_sync())
def test_create_new(self):
"""
Make sure that creating a new instance works.
"""
Manager.objects().create(name="Maz").run_sync()
names = [i["name"] for i in Manager.select(Manager.name).run_sync()]
self.assertTrue("Maz" in names)
async def run_transaction():
try:
async with Band._meta.db.transaction():
Manager.create_table()
Band.create_table()
Band.raw("MALFORMED QUERY ... SHOULD ERROR")
except Exception:
pass
def test_set_default(self):
Manager.alter().set_default(Manager.name, "Pending").run_sync()
# Bypassing the ORM to make sure the database default is present.
Band.raw("INSERT INTO manager (id, name) VALUES (DEFAULT, DEFAULT);"
).run_sync()
manager = Manager.objects().first().run_sync()
self.assertEqual(manager.name, "Pending")
def test_repr_postgres(self):
self.assertEqual(
Manager().__repr__(),
"<Manager: None>",
)
self.insert_row()
manager = Manager.objects().first().run_sync()
self.assertEqual(manager.__repr__(), "<Manager: 1>")
def test_succeeds(self):
transaction = Band._meta.db.atomic()
transaction.add(Manager.create_table(), Band.create_table())
transaction.run_sync()
self.assertTrue(Band.table_exists().run_sync())
self.assertTrue(Manager.table_exists().run_sync())
transaction.add(Band.alter().drop_table(),
Manager.alter().drop_table())
transaction.run_sync()
def test_drop_tables(self):
"""
Make sure the tables are dropped.
"""
self.assertEqual(Manager.table_exists().run_sync(), True)
self.assertEqual(Band.table_exists().run_sync(), True)
drop_tables(Manager, Band)
self.assertEqual(Manager.table_exists().run_sync(), False)
self.assertEqual(Band.table_exists().run_sync(), False)
def test_save_specific_columns(self):
"""
Make sure that we can save a subset of columns.
"""
manager = Manager(name="Guido")
manager.save().run_sync()
band = Band(name="Pythonistas", popularity=1000, manager=manager)
band.save().run_sync()
self.assertEqual(
Band.select().run_sync(),
[{
"id": 1,
"name": "Pythonistas",
"manager": 1,
"popularity": 1000,
}],
)
band.name = "Pythonistas 2"
band.popularity = 2000
band.save(columns=[Band.name]).run_sync()
# Only the name should update, and not the popularity:
self.assertEqual(
Band.select().run_sync(),
[{
"id": 1,
"name": "Pythonistas 2",
"manager": 1,
"popularity": 1000,
}],
)
#######################################################################
# Also test it using strings to identify columns
band.name = "Pythonistas 3"
band.popularity = 3000
band.save(columns=["popularity"]).run_sync()
# Only the popularity should update, and not the name:
self.assertEqual(
Band.select().run_sync(),
[{
"id": 1,
"name": "Pythonistas 2",
"manager": 1,
"popularity": 3000,
}],
)
def setUp(self):
for table in self.tables:
table.create_table().run_sync()
manager_1 = Manager(name="Guido")
manager_1.save().run_sync()
band_1 = Band(
name="Pythonistas", manager=manager_1.id, popularity=1000
)
band_1.save().run_sync()
manager_2 = Manager(name="Graydon")
manager_2.save().run_sync()
band_2 = Band(name="Rustaceans", manager=manager_2.id)
band_2.save().run_sync()
venue = Venue(name="Grand Central", capacity=1000)
venue.save().run_sync()
concert = Concert(band_1=band_1.id, band_2=band_2.id, venue=venue.id)
concert.save().run_sync()
ticket = Ticket(concert=concert, price=decimal.Decimal(50.0))
ticket.save().run_sync()
def test_where_equals(self):
self.insert_row()
manager = Manager.objects().first().run_sync()
# This is the recommended way of running these types of queries:
response = (Band.select(Band.name).where(Band.manager.id == getattr(
manager, Band._meta.primary_key._meta.name)).run_sync())
self.assertEqual(response, [{"name": "Pythonistas"}])
# Other cases which should work:
response = (Band.select(Band.name).where(Band.manager == getattr(
manager, Manager._meta.primary_key._meta.name)).run_sync())
self.assertEqual(response, [{"name": "Pythonistas"}])
response = (Band.select(Band.name).where(Band.manager.id == getattr(
manager, Manager._meta.primary_key._meta.name)).run_sync())
self.assertEqual(response, [{"name": "Pythonistas"}])
# check multiple arguments inside WHERE clause
response = (Band.select(Band.name).where(
Band.manager.id == 1, Band.popularity == 500).run_sync())
self.assertEqual(response, [])
# check empty WHERE clause
response = Band.select(Band.name).where().run_sync()
self.assertEqual(response, [{"name": "Pythonistas"}])
def test_error(self):
"""
Make sure queries in a transaction aren't committed if a query fails.
"""
transaction = Band._meta.db.atomic()
transaction.add(
Manager.create_table(),
Band.create_table(),
Band.raw("MALFORMED QUERY ... SHOULD ERROR"),
)
try:
transaction.run_sync()
except Exception:
pass
self.assertTrue(not Band.table_exists().run_sync())
self.assertTrue(not Manager.table_exists().run_sync())
def test_model_builder_with_valid_column_string(self):
manager = ModelBuilder.build_sync(Manager, defaults={"name": "Guido"})
queried_manager = (Manager.objects().where(
Manager.id == manager.id).first().run_sync())
self.assertEqual(queried_manager.name, "Guido")
def test_remove(self):
manager = Manager(name="Maz")
manager.save().run_sync()
self.assertTrue("Maz" in Manager.select(Manager.name).output(
as_list=True).run_sync())
manager.remove().run_sync()
self.assertTrue("Maz" not in Manager.select(Manager.name).output(
as_list=True).run_sync())
def setUp(self):
create_tables(*TABLES)
manager_1 = Manager.objects().create(name="Guido").run_sync()
manager_2 = Manager.objects().create(name="Graydon").run_sync()
band_1 = (Band.objects().create(name="Pythonistas",
manager=manager_1).run_sync())
band_2 = (Band.objects().create(name="Rustaceans",
manager=manager_2).run_sync())
venue = (Venue.objects().create(name="Royal Albert Hall",
capacity=5900).run_sync())
concert = (Concert.objects().create(venue=venue,
band_1=band_1,
band_2=band_2).run_sync())
Ticket.objects().create(price=decimal.Decimal(50.0),
concert=concert).run_sync()
def test_filter_rows(self):
"""
Make sure that `to_dict` works correctly with a subset of columns.
"""
self.insert_row()
instance = Manager.objects().first().run_sync()
dictionary = instance.to_dict(Manager.name)
self.assertDictEqual(dictionary, {"name": "Guido"})
def test_to_dict(self):
"""
Make sure that `to_dict` works correctly.
"""
self.insert_row()
instance = Manager.objects().first().run_sync()
dictionary = instance.to_dict()
self.assertDictEqual(dictionary, {"id": 1, "name": "Guido"})
def test_succeeds(self):
async def run_transaction():
async with Band._meta.db.transaction():
await Manager.create_table().run()
await Band.create_table().run()
asyncio.run(run_transaction())
self.assertTrue(Band.table_exists().run_sync())
self.assertTrue(Manager.table_exists().run_sync())
def test_str(self):
self.assertEqual(
Manager._table_str(),
(
"class Manager(Table, tablename='manager'):\n"
" id = Serial(null=False, primary_key=True, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name='id', secret=False)\n" # noqa: E501
" name = Varchar(length=50, default='', null=False, primary_key=False, unique=False, index=False, index_method=IndexMethod.btree, choices=None, db_column_name=None, secret=False)\n" # noqa: E501
),
)
self.assertEqual(
Manager._table_str(abbreviated=True),
("class Manager(Table):\n"
" id = Serial()\n"
" name = Varchar()\n"),
)
# We should also be able to print it directly.
print(Manager)
async def run_queries():
return await asyncio.gather(
Manager.select(),
Manager.insert(Manager(name="Golangs")),
Manager.delete().where(Manager.name != "Golangs"),
Manager.objects(),
Manager.count(),
Manager.raw("SELECT * FROM manager"),
)
def test_aliases(self):
"""
Make sure that `to_dict` works correctly with aliases.
"""
self.insert_row()
instance = Manager.objects().first().run_sync()
dictionary = instance.to_dict(
Manager.id, Manager.name.as_alias("title")
)
self.assertDictEqual(dictionary, {"id": 1, "title": "Guido"})
def test_get_related(self):
"""
Make sure you can get a related object from another object instance.
"""
manager = Manager(name="Guido")
manager.save().run_sync()
band = Band(name="Pythonistas", manager=manager.id, popularity=100)
band.save().run_sync()
_manager = band.get_related(Band.manager).run_sync()
self.assertTrue(_manager.name == "Guido")
# Test non-ForeignKey
with self.assertRaises(ValueError):
band.get_related(Band.name)
# Make sure it also works using a string
_manager = band.get_related("manager").run_sync()
self.assertTrue(_manager.name == "Guido")
# Test an invalid string
with self.assertRaises(ValueError):
band.get_related("abc123")
def test_alter_column_add_index(self):
"""
Test altering a column to add an index with MigrationManager.
"""
manager = MigrationManager()
manager.alter_column(
table_class_name="Manager",
tablename="manager",
column_name="name",
params={"index": True},
old_params={"index": False},
)
asyncio.run(manager.run())
self.assertTrue(
Manager._get_index_name(["name"]) in Manager.indexes().run_sync()
)
asyncio.run(manager.run_backwards())
self.assertTrue(
Manager._get_index_name(["name"])
not in Manager.indexes().run_sync()
)
def test_multiple(self):
self.insert_row()
query = (Manager.alter().add_column("column_a",
Integer(default=0,
null=True)).add_column(
"column_b",
Integer(default=0,
null=True)))
query.run_sync()
response = Band.raw("SELECT * FROM manager").run_sync()
column_names = response[0].keys()
self.assertIn("column_a", column_names)
self.assertIn("column_b", column_names)
def test_error(self):
"""
Make sure queries in a transaction aren't committed if a query fails.
"""
async def run_transaction():
try:
async with Band._meta.db.transaction():
Manager.create_table()
Band.create_table()
Band.raw("MALFORMED QUERY ... SHOULD ERROR")
except Exception:
pass
asyncio.run(run_transaction())
self.assertTrue(not Band.table_exists().run_sync())
self.assertTrue(not Manager.table_exists().run_sync())
def test_create_index(self):
"""
Test single column and multi column indexes.
"""
for columns in [
[Manager.name],
[Manager._meta.primary_key, Manager.name],
]:
Manager.create_index(columns).run_sync()
index_name = Manager._get_index_name(
[i._meta.name for i in columns])
index_names = Manager.indexes().run_sync()
self.assertIn(index_name, index_names)
Manager.drop_index(columns).run_sync()
index_names = Manager.indexes().run_sync()
self.assertNotIn(index_name, index_names)
def test_unique(self):
unique_query = Manager.alter().set_unique(Manager.name, True)
unique_query.run_sync()
Manager(name="Bob").save().run_sync()
# Make sure non-unique names work:
Manager(name="Sally").save().run_sync()
# Check there's a unique row error ...
with self.assertRaises(Exception):
Manager(name="Bob").save().run_sync()
response = Manager.select().run_sync()
self.assertEqual(len(response), 2)
# Now remove the constraint, and add a row.
not_unique_query = Manager.alter().set_unique(Manager.name, False)
not_unique_query.run_sync()
Manager(name="Bob").save().run_sync()
response = Manager.select().run_sync()
self.assertTrue(len(response), 2)
def test_save_new(self):
"""
Make sure that saving a new instance works.
"""
manager = Manager(name="Maz")
query = manager.save()
print(query)
self.assertTrue("INSERT" in query.__str__())
query.run_sync()
names = [i["name"] for i in Manager.select(Manager.name).run_sync()]
self.assertTrue("Maz" in names)
manager.name = "Maz2"
query = manager.save()
print(query)
self.assertTrue("UPDATE" in query.__str__())
query.run_sync()
names = [i["name"] for i in Manager.select(Manager.name).run_sync()]
self.assertTrue("Maz2" in names)
self.assertTrue("Maz" not in names)
def test_table_exists(self):
response = Manager.table_exists().run_sync()
self.assertTrue(response)
def setUp(self):
Manager.create_table().run_sync()
def tearDown(self):
Manager.alter().drop_table().run_sync()
