Skip to content
This repository has been archived by the owner on May 14, 2024. It is now read-only.

Commit

Permalink
Db new account insert (#308)
Browse files Browse the repository at this point in the history
* use similar validation of clusters to account logic

* Fix/refactor code that was trying to build the usage table with subtables

* clean up formatting on info output

* fix cartesian product warning output

* add sql expression form for Allocation.isExhausted hybrid property

* Add reminder to put in isExhausted tests

* Add function under Account Services to manually insert a db entry for a new account

* Add CLI option for Account insert

* Add tests for Account insert cli option

* Add test for Account insert account logic

* UpdateStatus is and AdminServices functionality, not AccountServices

* use subparsers

* fix accidentally breaking the usage table with merge resolution

* codacy items, missing setup for tests

* Caught trying to lazy load

* small fixes to tests

* Fix change from confusion of account update status and admin update status

* simplify account add

* fix assertion checking for finding an account

* missing EOF line

* fix handling empty usage_data

* Pull insert cli functionality in favor of automatic insertion

* use static method to perform account insertion if it does not already exist

* test automatic account insertion upon services object creation

* Add idempotence test on account setup in db

* typo

---------

Co-authored-by: Nickolas Comeau <[email protected]>
  • Loading branch information
Comeani and Comeani authored Aug 16, 2023
1 parent 3434e8d commit 0f0b306
Show file tree
Hide file tree
Showing 6 changed files with 112 additions and 5 deletions.
27 changes: 24 additions & 3 deletions bank/account_logic.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ def __init__(self, account_name: str) -> None:

account = SlurmAccount(account_name)
self._account_name = account.account_name
AccountServices.setup_db_account_entry(self._account_name)

def _get_active_proposal_id(self) -> None:
"""Return the active proposal ID for the current account
Expand Down Expand Up @@ -301,6 +302,7 @@ def __init__(self, account_name: str) -> None:

account = SlurmAccount(account_name)
self._account_name = account.account_name
AccountServices.setup_db_account_entry(self._account_name)

with DBConnection.session() as session:
# Check if the Account has an associated proposal
Expand Down Expand Up @@ -374,7 +376,7 @@ def create(
sus: The total number of service units to be added to the investment, or split equally across multiple
investments with ``num_inv``
start: The start date of the investment, or first in the sequence of investments, defaulting to today
end: The expiration date of the investment, or first in the sequence of investments,
end: The expiration date of the investment, or first in the sequence of investments,
defaulting to 12 months from ``start``
num_inv: Divide ``sus`` equally across a number of sequential investments
Expand Down Expand Up @@ -612,6 +614,7 @@ def __init__(self, account_name: str) -> None:

account = SlurmAccount(account_name)
self._account_name = account.account_name
self.setup_db_account_entry(self._account_name)

subquery = select(Account.id).where(Account.name == self._account_name)

Expand Down Expand Up @@ -671,11 +674,11 @@ def _build_usage_table(self) -> PrettyTable:
raise MissingProposalError('Account has no proposal')

# Proposal End Date as first row
output_table.add_row(['Proposal End Date:', proposal.end_date.strftime(settings.date_format),""],
output_table.add_row(['Proposal End Date:', proposal.end_date.strftime(settings.date_format), ""],
divider=True)

output_table.add_row(['Proposal ID:', proposal.id, ""], divider=True)
output_table.add_row(["","",""], divider=True)
output_table.add_row(["", "", ""], divider=True)

aggregate_usage_total = 0
allocation_total = 0
Expand All @@ -690,6 +693,8 @@ def _build_usage_table(self) -> PrettyTable:
continue

usage_data = slurm_acct.get_cluster_usage_per_user(allocation.cluster_name, in_hours=True)

# Skip if usage data is empty on the cluster
if not usage_data:
continue

Expand Down Expand Up @@ -793,6 +798,22 @@ def info(self) -> None:
except MissingInvestmentError:
pass

@staticmethod
def setup_db_account_entry(account_name) -> Account:
"""Insert an entry into the database for a new SLURM account if it does not already exist"""

with DBConnection.session() as session:
# Check if the Account has an entry in the database
account_query = select(Account).where(Account.name == account_name)
account = session.execute(account_query).scalars().first()

# If not, insert the new account so proposals/investments can reference it
if account is None:
account = Account(name=account_name)
session.add(account)
session.commit()
LOG.info(f"Created DB entry for account {account_name}")

def notify(self) -> None:
"""Send any pending usage alerts to the account"""

Expand Down
2 changes: 1 addition & 1 deletion bank/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ class CmdError(Exception):


class AccountNotFoundError(Exception):
"""Raised when a Slurm user account does not exist."""
"""Raised when a SLURM user account does not exist."""


class ClusterNotFoundError(Exception):
Expand Down
1 change: 1 addition & 0 deletions tests/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
.where(Investment.account_id.in_(account_subquery)) \
.where(Investment.is_active)


def add_proposal_to_test_account(proposal: Proposal) -> None:
"""Add a Proposal to the test account and commit the addition to the database """

Expand Down
50 changes: 49 additions & 1 deletion tests/account_logic/test_AccountServices.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,55 @@ def test_date_correct(self) -> None:
endDate = self.account._get_active_proposal_end_date()
self.assertEqual(endDate, date.today() + timedelta(days=365))


class SetupDBAccountEntry(TestCase):
"""Test first time insertion of the account into the DB"""

def test_account_inserted(self) -> None:
"""Test the account has an entry in the DB after insertion"""

account_name = settings.test_accounts[0]
# Insert an entry into the database for an account with an existing SLURM account
AccountServices.setup_db_account_entry(account_name)

with DBConnection.session() as session:
# Query the DB for the account
account_query = select(Account).where(Account.name == account_name)
account = session.execute(account_query).scalars().first()

# The account entry should not be empty, and the name should match the name provided
self.assertTrue(account)
self.assertEqual(account.name, settings.test_accounts[0])

def test_insertion_idempotence(self) -> None:
"""Test that multiple additions of the same account entry do not overwrite the initial insertion"""

account_name = settings.test_accounts[0]
AccountServices.setup_db_account_entry(account_name)
AccountServices.setup_db_account_entry(account_name)

with DBConnection.session() as session:
account_query = select(Account).where(Account.name == account_name)
accounts = session.execute(account_query).scalars().all()

self.assertEqual(len(accounts), 1)

def test_account_inserted_AccountServices(self) -> None:
"""Test the account has an entry in the DB upon AccountServices object creation"""

# Create an account services object for an existing SLURM account
acct = AccountServices(settings.test_accounts[0])

with DBConnection.session() as session:
# Query the DB for the account
account_query = select(Account).where(Account.name == acct._account_name)
account = session.execute(account_query).scalars().first()

# The account entry should not be empty, and the name should match the name provided
self.assertTrue(account)
self.assertEqual(account.name, acct._account_name)


@skip('This functionality hasn\'t been fully implemented yet.')
@patch('smtplib.SMTP.send_message')
class NotifyAccount(ProposalSetup, InvestmentSetup, TestCase):
Expand Down Expand Up @@ -355,4 +404,3 @@ def test_status_unlocked_with_investment_sus_applied(self) -> None:
investment = session.execute(active_investment_query).scalars().first()

self.assertEqual(900, investment.current_sus)

18 changes: 18 additions & 0 deletions tests/account_logic/test_InvestmentServices.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,24 @@ def test_error_on_negative_sus(self) -> None:
account.add_sus(1, 0)


class SetupDBAccountEntry(TestCase):
"""Test first time insertion of the account into the DB"""

def test_account_inserted_AccountServices(self) -> None:
"""Test the account has an entry in the DB upon InvestmentServices object creation"""

# Create an account services object for an existing SLURM account
acct = InvestmentServices(settings.test_accounts[0])

with DBConnection.session() as session:
# Query the DB for the account
account_query = select(Account).where(Account.name == acct._account_name)
account = session.execute(account_query).scalars().first()

# The account entry should not be empty, and the name should match the name provided
self.assertTrue(account)
self.assertEqual(account.name, acct._account_name)

class SubtractSus(ProposalSetup, InvestmentSetup, TestCase):
"""Tests for the subtraction of sus via the ``subtract`` method"""

Expand Down
19 changes: 19 additions & 0 deletions tests/account_logic/test_ProposalServices.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,25 @@ def test_error_on_negative_sus(self) -> None:
self.account.add_sus(**{settings.test_cluster: -1})


class SetupDBAccountEntry(TestCase):
"""Test first time insertion of the account into the DB"""

def test_account_inserted_AccountServices(self) -> None:
"""Test the account has an entry in the DB upon InvestmentServices object creation"""

# Create an account services object for an existing SLURM account
acct = ProposalServices(settings.test_accounts[0])

with DBConnection.session() as session:
# Query the DB for the account
account_query = select(Account).where(Account.name == acct._account_name)
account = session.execute(account_query).scalars().first()

# The account entry should not be empty, and the name should match the name provided
self.assertTrue(account)
self.assertEqual(account.name, acct._account_name)


class SubtractSus(ProposalSetup, TestCase):
"""Test the subtraction of sus via the ``subtract`` method"""

Expand Down

0 comments on commit 0f0b306

Please sign in to comment.