Super factories!

tests/domain/test_environments.py

@@ -16,6 +16,7 @@ from tests.factories import (
     WorkspaceFactory,
     EnvironmentFactory,
     ProjectFactory,
+    SuperWorkspaceFactory,
 )
 
 
@@ -55,20 +56,37 @@ def test_update_environment_roles():
 
 def test_get_scoped_environments(db):
     developer = UserFactory.create()
-    workspace = WorkspaceFactory.create()
-    workspace_user = Workspaces.add_member(workspace, developer, "developer")
-    project1 = ProjectFactory.create(workspace=workspace)
-    project2 = ProjectFactory.create(workspace=workspace)
-    env1 = EnvironmentFactory.create(project=project1, name="project1 dev")
-    env2 = EnvironmentFactory.create(project=project1, name="project1 staging")
-    env3 = EnvironmentFactory.create(project=project2, name="project2 dev")
-    env4 = EnvironmentFactory.create(project=project2, name="project2 staging")
-    db.session.add(EnvironmentRole(user=developer, environment=env1, role="developer"))
-    db.session.add(EnvironmentRole(user=developer, environment=env4, role="developer"))
-    db.session.commit()
+    workspace = SuperWorkspaceFactory.create(
+        name="hey",
+        members=[{"user": developer, "role_name": "developer"}],
+        projects=[
+            {
+                "name": "project1",
+                "environments": [
+                    {
+                        "name": "project1 dev",
+                        "members": [{"user": developer, "role_name": "developer"}],
+                    },
+                    {"name": "project1 staging"},
+                    {"name": "project1 prod"},
+                ],
+            },
+            {
+                "name": "project2",
+                "environments": [
+                    {"name": "project2 dev"},
+                    {
+                        "name": "project2 staging",
+                        "members": [{"user": developer, "role_name": "developer"}],
+                    },
+                    {"name": "project2 prod"},
+                ],
+            },
+        ],
+    )
 
-    project1_envs = Environments.for_user(developer, project1)
+    project1_envs = Environments.for_user(developer, workspace.projects[0])
     assert [env.name for env in project1_envs] == ["project1 dev"]
 
-    project2_envs = Environments.for_user(developer, project2)
+    project2_envs = Environments.for_user(developer, workspace.projects[1])
     assert [env.name for env in project2_envs] == ["project2 staging"]

tests/factories.py

@@ -18,6 +18,8 @@ from atst.models.user import User
 from atst.models.role import Role
 from atst.models.workspace import Workspace
 from atst.domain.roles import Roles
+from atst.models.workspace_role import WorkspaceRole
+from atst.models.environment_role import EnvironmentRole
 
 
 class Base(factory.alchemy.SQLAlchemyModelFactory):
@@ -98,6 +100,7 @@ class RequestFactory(Base):
         new_status=RequestStatus.STARTED,
         revision=factory.LazyAttribute(lambda se: se.factory_parent.revisions[-1]),
     )
+    task_order = factory.SubFactory("tests.factories.TaskOrderFactory")
 
     class Params:
         initial_revision = None
@@ -233,3 +236,84 @@ class ProjectFactory(Base):
 class EnvironmentFactory(Base):
     class Meta:
         model = Environment
+
+
+class WorkspaceRoleFactory(Base):
+    class Meta:
+        model = WorkspaceRole
+
+    workspace = factory.SubFactory(WorkspaceFactory)
+    role = factory.SubFactory(RoleFactory)
+    user = factory.SubFactory(UserFactory)
+
+
+class EnvironmentRoleFactory(Base):
+    class Meta:
+        model = EnvironmentRole
+
+    environment = factory.SubFactory(EnvironmentFactory)
+    role = factory.Faker("name")
+    user = factory.SubFactory(UserFactory)
+
+
+class SuperEnvironmentFactory(EnvironmentFactory):
+    @classmethod
+    def create(cls, *args, **kwargs):
+        with_members = kwargs.pop("members", [])
+        environment = super().create(*args, **kwargs)
+
+        for member in with_members:
+            user = member.get("user", UserFactory.create())
+            role_name = member["role_name"]
+            EnvironmentRoleFactory.create(
+                environment=environment, role=role_name, user=user
+            )
+
+        return environment
+
+
+class SuperProjectFactory(ProjectFactory):
+    @classmethod
+    def create(cls, *args, **kwargs):
+        with_environments = kwargs.pop("environments", [])
+        project = super().create(*args, **kwargs)
+
+        environments = [
+            SuperEnvironmentFactory.create(project=project, **e) for e in with_environments
+        ]
+
+        project.environments = environments
+        return project
+
+
+class SuperWorkspaceFactory(WorkspaceFactory):
+    class Meta:
+        pass
+
+    @classmethod
+    def create(cls, *args, **kwargs):
+        with_projects = kwargs.pop("projects", [])
+        owner = kwargs.pop("owner", None)
+        members = kwargs.pop("members", [])
+
+        workspace = super().create(*args, **kwargs)
+
+        projects = [
+            SuperProjectFactory.create(workspace=workspace, **p) for p in with_projects
+        ]
+
+        if owner:
+            workspace.request.creator = owner
+            WorkspaceRoleFactory.create(
+                workspace=workspace, role=Roles.get("owner"), user=owner
+            )
+
+        for member in members:
+            user = member.get("user", UserFactory.create())
+            role_name = member["role_name"]
+            WorkspaceRoleFactory.create(
+                workspace=workspace, role=Roles.get(role_name), user=user
+            )
+
+        workspace.projects = projects
+        return workspace