Currently, we use both Python's built-in datetime library and Pendulum to do datetime operations. For the sake of consistency, we should try to stick to one library for datetimes. We could have used either, but Pendulum has a more ergonomic API, so I decided to go with it when possible. The places where were we didn't / couldn't replace datetime are: - checking instances of datetimes. Pendulum's objects are subclasses of python native datetime objects, so it's still useful to import datetime in those cases of using is_instance() - WTForms date validators expect datetime style string formats -- Pendulum has its own format for formatting/ parsing strings. As such, our custom validator DateRange needs to use datetime.stptime() to account for this format.
378 lines
11 KiB
Python
378 lines
11 KiB
Python
# Add root application dir to the python path
|
|
import os
|
|
import sys
|
|
import pendulum
|
|
import random
|
|
from faker import Faker
|
|
from werkzeug.datastructures import FileStorage
|
|
from uuid import uuid4
|
|
|
|
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
sys.path.append(parent_dir)
|
|
|
|
from atst.app import make_config, make_app
|
|
from atst.database import db
|
|
|
|
from atst.models.application import Application
|
|
from atst.models.clin import JEDICLINType
|
|
from atst.models.environment_role import CSPRole
|
|
|
|
from atst.domain.application_roles import ApplicationRoles
|
|
from atst.domain.applications import Applications
|
|
from atst.domain.csp.reports import MockReportingProvider
|
|
from atst.domain.environments import Environments
|
|
from atst.domain.environment_roles import EnvironmentRoles
|
|
from atst.domain.exceptions import AlreadyExistsError, NotFoundError
|
|
from atst.domain.invitations import ApplicationInvitations
|
|
from atst.domain.permission_sets import PermissionSets, APPLICATION_PERMISSION_SETS
|
|
from atst.domain.portfolio_roles import PortfolioRoles
|
|
from atst.domain.portfolios import Portfolios
|
|
from atst.domain.users import Users
|
|
|
|
from atst.routes.dev import _DEV_USERS as DEV_USERS
|
|
|
|
from atst.utils import pick
|
|
|
|
from tests.factories import (
|
|
random_defense_component,
|
|
TaskOrderFactory,
|
|
CLINFactory,
|
|
AttachmentFactory,
|
|
)
|
|
|
|
fake = Faker()
|
|
|
|
|
|
PORTFOLIO_USERS = [
|
|
{
|
|
"first_name": "Danny",
|
|
"last_name": "Knight",
|
|
"email": "knight@mil.gov",
|
|
"dod_id": "0000000001",
|
|
"permission_sets": PortfolioRoles.DEFAULT_PORTFOLIO_PERMISSION_SETS,
|
|
},
|
|
{
|
|
"first_name": "Mario",
|
|
"last_name": "Hudson",
|
|
"email": "hudson@mil.gov",
|
|
"dod_id": "0000000002",
|
|
"permission_sets": PortfolioRoles.DEFAULT_PORTFOLIO_PERMISSION_SETS,
|
|
},
|
|
{
|
|
"first_name": "Louise",
|
|
"last_name": "Greer",
|
|
"email": "greer@mil.gov",
|
|
"dod_id": "0000000003",
|
|
"permission_sets": PortfolioRoles.DEFAULT_PORTFOLIO_PERMISSION_SETS,
|
|
},
|
|
]
|
|
|
|
|
|
APPLICATION_USERS = [
|
|
{
|
|
"first_name": "Jean Luc",
|
|
"last_name": "Picard",
|
|
"email": "picard@mil.gov",
|
|
"dod_id": "0000000004",
|
|
"permission_sets": random.sample(
|
|
APPLICATION_PERMISSION_SETS, k=random.randint(1, 4)
|
|
),
|
|
},
|
|
{
|
|
"first_name": "()",
|
|
"last_name": "Spock",
|
|
"email": "spock@mil.gov",
|
|
"dod_id": "0000000005",
|
|
"permission_sets": random.sample(
|
|
APPLICATION_PERMISSION_SETS, k=random.randint(1, 4)
|
|
),
|
|
},
|
|
{
|
|
"first_name": "William",
|
|
"last_name": "Shatner",
|
|
"email": "shatner@mil.gov",
|
|
"dod_id": "0000000006",
|
|
"permission_sets": random.sample(
|
|
APPLICATION_PERMISSION_SETS, k=random.randint(1, 4)
|
|
),
|
|
},
|
|
{
|
|
"first_name": "Nyota",
|
|
"last_name": "Uhura",
|
|
"email": "uhura@mil.gov",
|
|
"dod_id": "0000000007",
|
|
"permission_sets": random.sample(
|
|
APPLICATION_PERMISSION_SETS, k=random.randint(1, 4)
|
|
),
|
|
},
|
|
{
|
|
"first_name": "Kathryn",
|
|
"last_name": "Janeway",
|
|
"email": "janeway@mil.gov",
|
|
"dod_id": "0000000008",
|
|
"permission_sets": random.sample(
|
|
APPLICATION_PERMISSION_SETS, k=random.randint(1, 4)
|
|
),
|
|
},
|
|
]
|
|
|
|
|
|
SHIP_NAMES = [
|
|
"Millenium Falcon",
|
|
"Star Destroyer",
|
|
"Attack Cruiser",
|
|
"Sith Infiltrator",
|
|
"Death Star",
|
|
"Lambda Shuttle",
|
|
"Corellian Corvette",
|
|
]
|
|
|
|
|
|
SOFTWARE_WORDS = [
|
|
"Enterprise",
|
|
"Scalable",
|
|
"Solution",
|
|
"Blockchain",
|
|
"Cloud",
|
|
"Micro",
|
|
"Macro",
|
|
"Software",
|
|
"Global",
|
|
"Team",
|
|
]
|
|
|
|
|
|
ENVIRONMENT_NAMES = ["production", "staging", "test", "uat", "dev", "qa"]
|
|
|
|
|
|
def get_users():
|
|
users = []
|
|
for dev_user in DEV_USERS.values():
|
|
try:
|
|
user = Users.create(**dev_user)
|
|
except AlreadyExistsError:
|
|
user = Users.get_by_dod_id(dev_user["dod_id"])
|
|
|
|
users.append(user)
|
|
return users
|
|
|
|
|
|
def add_members_to_portfolio(portfolio):
|
|
for user_data in PORTFOLIO_USERS:
|
|
invite = Portfolios.invite(portfolio, portfolio.owner, {"user_data": user_data})
|
|
profile = {
|
|
k: user_data[k] for k in user_data if k not in ["dod_id", "permission_sets"]
|
|
}
|
|
user = Users.get_or_create_by_dod_id(user_data["dod_id"], **profile)
|
|
PortfolioRoles.enable(invite.role, user)
|
|
|
|
db.session.commit()
|
|
|
|
|
|
def add_task_orders_to_portfolio(portfolio):
|
|
today = pendulum.today()
|
|
future = today.add(days=100)
|
|
yesterday = today.subtract(days=1)
|
|
|
|
def build_pdf():
|
|
return {"filename": "sample_task_order.pdf", "object_name": str(uuid4())}
|
|
|
|
draft_to = TaskOrderFactory.build(portfolio=portfolio, pdf=None)
|
|
unsigned_to = TaskOrderFactory.build(portfolio=portfolio, pdf=build_pdf())
|
|
upcoming_to = TaskOrderFactory.build(
|
|
portfolio=portfolio, signed_at=yesterday, pdf=build_pdf()
|
|
)
|
|
expired_to = TaskOrderFactory.build(
|
|
portfolio=portfolio, signed_at=yesterday, pdf=build_pdf()
|
|
)
|
|
active_to = TaskOrderFactory.build(
|
|
portfolio=portfolio, signed_at=yesterday, pdf=build_pdf()
|
|
)
|
|
|
|
clins = [
|
|
CLINFactory.build(
|
|
task_order=unsigned_to, start_date=today.subtract(days=5), end_date=today
|
|
),
|
|
CLINFactory.build(
|
|
task_order=upcoming_to, start_date=today.add(days=5), end_date=future
|
|
),
|
|
CLINFactory.build(
|
|
task_order=expired_to, start_date=today.subtract(days=5), end_date=yesterday
|
|
),
|
|
CLINFactory.build(
|
|
task_order=active_to,
|
|
start_date=yesterday,
|
|
end_date=future,
|
|
total_amount=1_000_000,
|
|
obligated_amount=500_000,
|
|
jedi_clin_type=JEDICLINType.JEDI_CLIN_1,
|
|
),
|
|
CLINFactory.build(
|
|
task_order=active_to,
|
|
start_date=yesterday,
|
|
end_date=future,
|
|
total_amount=500_000,
|
|
obligated_amount=200_000,
|
|
jedi_clin_type=JEDICLINType.JEDI_CLIN_2,
|
|
),
|
|
]
|
|
|
|
task_orders = [draft_to, unsigned_to, upcoming_to, expired_to, active_to]
|
|
|
|
db.session.add_all(task_orders + clins)
|
|
db.session.commit()
|
|
|
|
|
|
def random_applications():
|
|
return [
|
|
{
|
|
"name": fake.sentence(nb_words=3, ext_word_list=SOFTWARE_WORDS)[0:-1],
|
|
"description": fake.bs(),
|
|
"environments": random.sample(ENVIRONMENT_NAMES, k=random.randint(1, 4)),
|
|
}
|
|
for n in range(random.randint(1, 4))
|
|
]
|
|
|
|
|
|
def add_applications_to_portfolio(portfolio):
|
|
applications = random_applications()
|
|
for application_data in applications:
|
|
application = Applications.create(
|
|
portfolio.owner,
|
|
portfolio=portfolio,
|
|
name=application_data["name"],
|
|
description=application_data["description"],
|
|
environment_names=application_data["environments"],
|
|
)
|
|
|
|
users = random.sample(APPLICATION_USERS, k=random.randint(1, 5))
|
|
for user_data in users:
|
|
try:
|
|
user = Users.get_by_dod_id(user_data["dod_id"])
|
|
except NotFoundError:
|
|
user = Users.create(
|
|
user_data["dod_id"],
|
|
None,
|
|
first_name=user_data["first_name"],
|
|
last_name=user_data["last_name"],
|
|
email=user_data["email"],
|
|
)
|
|
|
|
app_role = ApplicationRoles.create(
|
|
user=user,
|
|
application=application,
|
|
permission_set_names=[PermissionSets.EDIT_APPLICATION_TEAM],
|
|
)
|
|
|
|
ApplicationInvitations.create(
|
|
portfolio.owner, app_role, user_data, commit=True
|
|
)
|
|
|
|
user_environments = random.sample(
|
|
application.environments,
|
|
k=random.randint(1, len(application.environments)),
|
|
)
|
|
for env in user_environments:
|
|
role = random.choice([e.value for e in CSPRole])
|
|
EnvironmentRoles.create(
|
|
application_role=app_role, environment=env, role=role
|
|
)
|
|
|
|
|
|
def create_demo_portfolio(name, data):
|
|
try:
|
|
portfolio_owner = Users.get_or_create_by_dod_id(
|
|
"2345678901",
|
|
**pick(
|
|
[
|
|
"permission_sets",
|
|
"first_name",
|
|
"last_name",
|
|
"email",
|
|
"service_branch",
|
|
"phone_number",
|
|
"citizenship",
|
|
"designation",
|
|
"date_latest_training",
|
|
],
|
|
DEV_USERS["amanda"],
|
|
),
|
|
) # Amanda
|
|
# auditor = Users.get_by_dod_id("3453453453") # Sally
|
|
except NotFoundError:
|
|
print(
|
|
"Could not find demo users; will not create demo portfolio {}".format(name)
|
|
)
|
|
return
|
|
|
|
portfolio = Portfolios.create(
|
|
user=portfolio_owner,
|
|
portfolio_attrs={"name": name, "defense_component": random_defense_component()},
|
|
)
|
|
|
|
add_task_orders_to_portfolio(portfolio)
|
|
add_members_to_portfolio(portfolio)
|
|
|
|
for mock_application in data["applications"]:
|
|
application = Application(
|
|
portfolio=portfolio, name=mock_application["name"], description=""
|
|
)
|
|
env_names = [env["name"] for env in mock_application["environments"]]
|
|
envs = Environments.create_many(portfolio.owner, application, env_names)
|
|
db.session.add(application)
|
|
db.session.commit()
|
|
|
|
|
|
def seed_db():
|
|
get_users()
|
|
amanda = Users.get_by_dod_id("2345678901")
|
|
|
|
# Create Portfolios for Amanda with mocked reporting data
|
|
create_demo_portfolio("A-Wing", MockReportingProvider.FIXTURE_SPEND_DATA["A-Wing"])
|
|
create_demo_portfolio("B-Wing", MockReportingProvider.FIXTURE_SPEND_DATA["B-Wing"])
|
|
|
|
tie_interceptor = Portfolios.create(
|
|
user=amanda,
|
|
portfolio_attrs={
|
|
"name": "TIE Interceptor",
|
|
"defense_component": random_defense_component(),
|
|
},
|
|
)
|
|
add_task_orders_to_portfolio(tie_interceptor)
|
|
add_members_to_portfolio(tie_interceptor)
|
|
add_applications_to_portfolio(tie_interceptor)
|
|
|
|
tie_fighter = Portfolios.create(
|
|
user=amanda,
|
|
portfolio_attrs={
|
|
"name": "TIE Fighter",
|
|
"defense_component": random_defense_component(),
|
|
},
|
|
)
|
|
add_task_orders_to_portfolio(tie_fighter)
|
|
add_members_to_portfolio(tie_fighter)
|
|
add_applications_to_portfolio(tie_fighter)
|
|
|
|
# create a portfolio for each user
|
|
ships = SHIP_NAMES.copy()
|
|
for user in get_users():
|
|
ship = random.choice(ships)
|
|
ships.remove(ship)
|
|
portfolio = Portfolios.create(
|
|
user=user,
|
|
portfolio_attrs={
|
|
"name": ship,
|
|
"defense_component": random_defense_component(),
|
|
},
|
|
)
|
|
add_task_orders_to_portfolio(portfolio)
|
|
add_members_to_portfolio(portfolio)
|
|
add_applications_to_portfolio(portfolio)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
config = make_config({"DISABLE_CRL_CHECK": True, "DEBUG": False})
|
|
app = make_app(config)
|
|
with app.app_context():
|
|
seed_db()
|