mage-ai 0.9.68__py3-none-any.whl → 0.9.70__py3-none-any.whl

mage_ai/server/server.py

@@ -13,6 +13,7 @@ from typing import Optional, Union
 import pytz
 import tornado.ioloop
 import tornado.web
+from sqlalchemy import or_
 from tornado import autoreload
 from tornado.ioloop import PeriodicCallback
 from tornado.log import enable_pretty_logging
@@ -42,7 +43,7 @@ from mage_ai.data_preparation.repo_manager import (
 )
 from mage_ai.data_preparation.shared.constants import MANAGE_ENV_VAR
 from mage_ai.orchestration.constants import Entity
-from mage_ai.orchestration.db import db_connection, set_db_schema
+from mage_ai.orchestration.db import db_connection, safe_db_query, set_db_schema
 from mage_ai.orchestration.db.database_manager import database_manager
 from mage_ai.orchestration.db.models.oauth import Oauth2Application, Role, User
 from mage_ai.orchestration.utils.distributed_lock import DistributedLock
@@ -55,6 +56,7 @@ from mage_ai.server.api.events import (
     ApiEventMatcherDetailHandler,
     ApiEventMatcherListHandler,
 )
+from mage_ai.server.api.runs import ApiRunHandler
 from mage_ai.server.api.triggers import ApiTriggerPipelineHandler
 from mage_ai.server.api.v1 import (
     ApiChildDetailHandler,
@@ -321,6 +323,16 @@ def make_app(
             ApiTriggerPipelineHandler,
         ),
 
+        # Run a single block and get a response immediately
+        (
+            r'/api/runs',
+            ApiRunHandler,
+        ),
+        (
+            r'/api/runs/(?P<token>\w+)',
+            ApiRunHandler,
+        ),
+
         # Download block output
         (
             r'/api/pipelines/(?P<pipeline_uuid>\w+)/block_outputs/'
@@ -444,6 +456,7 @@ def make_app(
     )
 
 
+@safe_db_query
 def initialize_user_authentication(project_type: ProjectType) -> Oauth2Application:
     logger.info('User authentication is enabled.')
     # We need to sleep for a few seconds after creating all the tables or else there
@@ -452,12 +465,14 @@ def initialize_user_authentication(project_type: ProjectType) -> Oauth2Applicati
 
     # Create new roles on existing users. This should only need to be run once.
     if project_type == ProjectType.SUB:
+        project_uuid = get_project_uuid()
+        project_uuid_truncated = project_uuid[:8]
         Role.create_default_roles(
             entity=Entity.PROJECT,
-            entity_id=get_project_uuid(),
-            prefix=get_repo_name(),
+            entity_id=project_uuid,
+            name_func=lambda role: f'{role}_{project_uuid_truncated}',
         )
-        default_owner_role = Role.get_role(f'{get_repo_name()}_{Role.DefaultRole.OWNER}')
+        default_owner_role = Role.get_role(f'{Role.DefaultRole.OWNER}_{project_uuid_truncated}')
     else:
         Role.create_default_roles()
         default_owner_role = Role.get_role(Role.DefaultRole.OWNER)
@@ -469,19 +484,29 @@ def initialize_user_authentication(project_type: ProjectType) -> Oauth2Applicati
     if not legacy_owner_user and len(owner_users) == 0:
         logger.info('User with owner permission doesn’t exist, creating owner user.')
         if AUTHENTICATION_MODE.lower() == 'ldap':
-            user = User.create(
-                roles_new=[default_owner_role],
-                username=get_settings_value(LDAP_ADMIN_USERNAME, 'admin'),
-            )
+            username = get_settings_value(LDAP_ADMIN_USERNAME, 'admin')
+            user = User.query.filter(User.username == username).first()
+            if not user:
+                user = User.create(
+                    roles_new=[default_owner_role],
+                    username=get_settings_value(LDAP_ADMIN_USERNAME, 'admin'),
+                )
         else:
             password_salt = generate_salt()
-            user = User.create(
-                email='admin@admin.com',
-                password_hash=create_bcrypt_hash('admin', password_salt),
-                password_salt=password_salt,
-                roles_new=[default_owner_role],
-                username='admin',
-            )
+            user = User.query.filter(
+                or_(
+                    User.email == 'admin@admin.com',
+                    User.username == 'admin',
+                ),
+            ).first()
+            if not user:
+                user = User.create(
+                    email='admin@admin.com',
+                    password_hash=create_bcrypt_hash('admin', password_salt),
+                    password_salt=password_salt,
+                    roles_new=[default_owner_role],
+                    username='admin',
+                )
         owner_user = user
     else:
         if legacy_owner_user and not legacy_owner_user.roles_new:
@@ -573,7 +598,7 @@ async def main(
             from mage_ai.data_preparation.sync.git_sync import GitSync
             sync_config = GitConfig.load(config=preferences.sync_config)
             sync = GitSync(sync_config, setup_repo=True)
-            if sync_config.sync_on_start:
+            if sync_config.remote_repo_link and sync_config.sync_on_start is True:
                 try:
                     sync.sync_data()
                     logger.info(
@@ -671,7 +696,12 @@ async def main(
     update_settings_on_metadata_change()
     observer = Observer()
     event_handler = MetadataEventHandler()
-    observer.schedule(event_handler, path=get_metadata_path(root_project=True))
+    metadata_file = get_metadata_path(root_project=True)
+    if not os.path.exists(metadata_file):
+        os.makedirs(os.path.dirname(metadata_file), exist_ok=True)
+        with open(metadata_file, 'w') as f:
+            f.write('')
+    observer.schedule(event_handler, path=metadata_file)
     observer.start()
 
     get_messages(

mage_ai/server/utils/output_display.py

@@ -208,7 +208,7 @@ def __custom_output():
             ignore_nan=True,
         )
         return print(f'[__internal_output__]{{_json_string}}')
-    elif bool({has_reduce_output}):
+    elif bool({has_reduce_output}) and bool({is_dynamic_child}):
         _json_string = simplejson.dumps(
             transform_output_for_display_reduce_output(
                 _internal_output_return,
@@ -460,7 +460,7 @@ def execute_custom_code():
 
     output = block_output['output'] or []
 
-    if {widget} or is_dynamic_block(block) or is_dynamic_child or has_reduce_output:
+    if {widget} or is_dynamic_block(block) or is_dynamic_child:
         return output
     else:
         return find(lambda val: val is not None, output)

mage_ai/server/websocket_server.py

@@ -116,6 +116,7 @@ def run_pipeline(
         pipeline.execute_sync(
             global_vars=global_vars,
             build_block_output_stdout=build_block_output_stdout,
+            retry_config=dict(),
             run_sensors=False,
         )
         add_pipeline_message(

mage_ai/services/aws/ecs/ecs.py

@@ -16,6 +16,7 @@ def run_task(
     response = client.run_task(**ecs_config.get_task_config(command=command))
 
     print(json.dumps(response, indent=4, default=str))
+    wait_for_completion = False if ecs_config.wait_timeout == -1 else wait_for_completion
 
     if wait_for_completion:
         arn = response['tasks'][0]['taskArn']

mage_ai/services/k8s/config.py

@@ -2,7 +2,6 @@ from dataclasses import dataclass
 from typing import Dict
 
 from kubernetes.client import (
-    V1Affinity,
     V1Container,
     V1EnvVar,
     V1LocalObjectReference,
@@ -15,8 +14,9 @@ from kubernetes.client import (
 )
 
 from mage_ai.services.k8s.constants import CONFIG_FILE, DEFAULT_NAMESPACE
+from mage_ai.services.k8s.utils import parse_affinity_config
 from mage_ai.shared.config import BaseConfig
-from mage_ai.shared.hash import camel_case_keys_to_snake_case, get_safe_value
+from mage_ai.shared.hash import get_safe_value
 
 # import traceback
 
@@ -61,8 +61,8 @@ class K8sExecutorConfig(BaseConfig):
                 executor_config.pod.get('service_account_name') or DEFAULT_SERVICE_ACCOUNT_NAME
             )
             if executor_config.pod.get('affinity'):
-                affinity = V1Affinity(
-                    **camel_case_keys_to_snake_case(executor_config.pod['affinity']))
+                # Convert the affinity to a V1Affinity object
+                affinity = parse_affinity_config(executor_config.pod['affinity'])
 
             if executor_config.pod.get('tolerations'):
                 tolerations += [V1Toleration(**e) for e in executor_config.pod['tolerations']]

mage_ai/services/k8s/utils.py

@@ -0,0 +1,97 @@
+from kubernetes import client
+
+
+def parse_affinity_config(affinity_config):
+    if not affinity_config:
+        return None
+
+    affinity = client.V1Affinity(
+        node_affinity=parse_node_affinity(affinity_config.get('nodeAffinity')),
+        pod_affinity=parse_pod_affinity(affinity_config.get('podAffinity')),
+        pod_anti_affinity=parse_pod_anti_affinity(affinity_config.get('podAntiAffinity'))
+    )
+    return affinity
+
+
+def parse_node_affinity(node_affinity_config):
+    if not node_affinity_config:
+        return None
+
+    required_terms = node_affinity_config.get(
+        'requiredDuringSchedulingIgnoredDuringExecution', {}).get('nodeSelectorTerms', [])
+    node_selector_terms = []
+    for term in required_terms:
+        match_expressions = [client.V1NodeSelectorRequirement(
+            key=expr.get('key'),
+            operator=expr.get('operator'),
+            values=expr.get('values')
+        ) for expr in term.get('matchExpressions', [])]
+
+        node_selector_terms.append(client.V1NodeSelectorTerm(match_expressions=match_expressions))
+
+    return client.V1NodeAffinity(
+        required_during_scheduling_ignored_during_execution=client.V1NodeSelector(
+            node_selector_terms=node_selector_terms)
+    )
+
+
+def parse_pod_affinity(pod_affinity_config):
+    if not pod_affinity_config:
+        return None
+
+    required_terms = pod_affinity_config.get('requiredDuringSchedulingIgnoredDuringExecution', [])
+    preferred_terms = pod_affinity_config.get('preferredDuringSchedulingIgnoredDuringExecution', [])
+
+    required_affinity = parse_affinity_term_list(required_terms)
+    preferred_affinity = parse_affinity_term_list(preferred_terms)
+
+    return client.V1PodAffinity(
+        required_during_scheduling_ignored_during_execution=required_affinity,
+        preferred_during_scheduling_ignored_during_execution=preferred_affinity
+    )
+
+
+def parse_pod_anti_affinity(pod_anti_affinity_config):
+    if not pod_anti_affinity_config:
+        return None
+
+    required_terms = pod_anti_affinity_config.get(
+        'requiredDuringSchedulingIgnoredDuringExecution', [])
+    preferred_terms = pod_anti_affinity_config.get(
+        'preferredDuringSchedulingIgnoredDuringExecution', [])
+
+    required_affinity = parse_affinity_term_list(required_terms)
+    preferred_affinity = parse_affinity_term_list(preferred_terms)
+
+    return client.V1PodAntiAffinity(
+        required_during_scheduling_ignored_during_execution=required_affinity,
+        preferred_during_scheduling_ignored_during_execution=preferred_affinity
+    )
+
+
+def parse_affinity_term_list(term_list):
+    if not term_list:
+        return None
+
+    affinity_term_list = []
+    for term in term_list:
+        label_selector = client.V1LabelSelector(
+            match_labels=term.get('labelSelector', {}).get('matchLabels', {}),
+            match_expressions=[client.V1LabelSelectorRequirement(
+                key=expr.get('key'),
+                operator=expr.get('operator'),
+                values=expr.get('values')
+            ) for expr in term.get('labelSelector', {}).get('matchExpressions', [])]
+        )
+
+        topology_key = term.get('topologyKey')
+
+        affinity_term_list.append(client.V1WeightedPodAffinityTerm(
+            weight=term.get('weight'),
+            pod_affinity_term=client.V1PodAffinityTerm(
+                label_selector=label_selector,
+                topology_key=topology_key
+            )
+        ))
+
+    return affinity_term_list

mage_ai/settings/keys/auth.py

@@ -13,6 +13,8 @@ LDAP_DEFAULT_ACCESS = 'LDAP_DEFAULT_ACCESS'
 LDAP_GROUP_FIELD = 'LDAP_GROUP_FIELD'
 LDAP_ROLES_MAPPING = 'LDAP_ROLES_MAPPING'
 
+UPDATE_ROLES_ON_LOGIN = 'UPDATE_ROLES_ON_LOGIN'
+
 # Git settings
 
 GIT_REPO_LINK = 'GIT_REPO_LINK'

mage_ai/settings/server.py

@@ -79,7 +79,7 @@ SERVER_LOGGING_TEMPLATE = os.getenv(
     '%(levelname)s:%(name)s:%(message)s',
 )
 
-INITIAL_METADATA = os.getenv('INITIAL_METADATA', '{{}}')
+INITIAL_METADATA = os.getenv('INITIAL_METADATA')
 
 DISABLE_AUTO_BROWSER_OPEN = get_bool_value(os.getenv('DISABLE_AUTO_BROWSER_OPEN', 'False'))
 

mage_ai/shared/parsers.py

@@ -25,11 +25,13 @@ MAX_ITEMS_IN_SAMPLE_OUTPUT = 20
 
 
 def encode_complex(obj):
+    from mage_ai.shared.models import BaseDataClass
+
     if isinstance(obj, set):
         return list(obj)
     elif isinstance(obj, BaseModel):
         return obj.__class__.__name__
-    elif obj.__class__.__name__ == 'BaseDataClass':
+    elif obj.__class__.__name__ == 'BaseDataClass' or isinstance(obj, BaseDataClass):
         return obj.to_dict()
     elif isinstance(obj, Enum):
         return obj.value
@@ -59,6 +61,9 @@ def encode_complex(obj):
         return obj.to_dict(orient='records')
     elif isinstance(obj, pd.Series):
         return obj.to_list()
+    # Convert pandas._libs.missing.NAType to None
+    elif isinstance(obj, pd._libs.missing.NAType):
+        return None
 
     return obj
 

mage_ai/streaming/sources/influxdb.py

@@ -164,6 +164,7 @@ class InfluxDbSource(BaseSource):
                         'data': {record.get_field(): record.get_value()},
                         'metadata': {
                             'time': int(1e3 * record.get_time().timestamp()),
+                            'measurement': record['_measurement'],
                             'tags': {
                                 k: v
                                 for k, v in record.values.items()
@@ -201,6 +202,7 @@ class InfluxDbSource(BaseSource):
                         'data': {record.get_field(): record.get_value()},
                         'metadata': {
                             'time': int(1e3 * record.get_time().timestamp()),
+                            'measurement': record['_measurement'],
                             'tags': {
                                 k: v
                                 for k, v in record.values.items()

mage_ai/streaming/sources/kafka.py

@@ -257,7 +257,7 @@ class KafkaSource(BaseSource):
 
     def test_connection(self):
         self.consumer._client.check_version(timeout=5)
-        self._print('Test connectino successfully.')
+        self._print('Test connection successfully.')
 
     def __deserialize_message(self, message):
         if self.config.serde_config is None:

mage_ai/tests/api/endpoints/mixins.py

@@ -813,7 +813,7 @@ class BaseAPIEndpointTest(AsyncDBTestCase):
                             validation = assert_after_create_count(self)
                         self.assertTrue(validation)
                     else:
-                        after_count = len(get_resource().model_class.all())
+                        after_count = len(get_resource(resource).model_class.all())
                         self.assertEqual(
                             after_count,
                             before_count + 1 if after_create_count is None else after_create_count,
@@ -1096,7 +1096,7 @@ class BaseAPIEndpointTest(AsyncDBTestCase):
 
                         self.assertTrue(validation)
                     else:
-                        after_count = len(get_resource().model_class.all())
+                        after_count = len(get_resource(resource).model_class.all())
                         self.assertEqual(
                             after_count,
                             before_count - 1 if after_delete_count is None else after_delete_count,
@@ -1132,15 +1132,16 @@ class BaseAPIEndpointTest(AsyncDBTestCase):
                 )
                 RolePermission.create(permission_id=permission.id, role_id=role.id)
             else:
-                permission = Permission.create(
-                    access=access_for_permissions,
-                    entity_name=EntityName(entity_name(resource)),
-                    options=permission_options,
-                )
+                e_name = entity_name(resource)
+                if getattr(EntityName, e_name, None) is not None:
+                    permission = Permission.create(
+                        access=access_for_permissions,
+                        entity_name=EntityName(e_name),
+                        options=permission_options,
+                    )
+                    RolePermission.create(permission_id=permission.id, role_id=role.id)
 
                 if permission_settings:
                     for permission_setting in permission_settings:
                         permission_more = Permission.create(**permission_setting)
                         RolePermission.create(permission_id=permission_more.id, role_id=role.id)
-
-            RolePermission.create(permission_id=permission.id, role_id=role.id)

mage_ai/tests/api/endpoints/test_seeds.py

@@ -0,0 +1,24 @@
+from mage_ai.authentication.permissions.constants import PermissionAccess
+from mage_ai.tests.api.endpoints.mixins import (
+    BaseAPIEndpointTest,
+    build_create_endpoint_tests,
+)
+
+
+class SeedAPIEndpointTest(BaseAPIEndpointTest):
+    pass
+
+
+build_create_endpoint_tests(
+    SeedAPIEndpointTest,
+    resource='seed',
+    authentication_accesses=[PermissionAccess.OWNER],
+    permissions_accesses=[PermissionAccess.OWNER],
+    build_payload=lambda self: dict(
+        roles=True,
+        permissions=True,
+        policy_names=['Block', 'Pipeline', 'Download'],
+    ),
+    assert_before_create_count=lambda self: True,
+    assert_after_create_count=lambda self: True,
+)

mage_ai/tests/api/operations/base/mixins.py

@@ -34,7 +34,7 @@ class BaseOperationMixin:
                 self.faker.unique.random_int(),
             ],
         )
-        self.user = User(username=self.faker.unique.name())
+        self.user = User.create(username=self.faker.unique.name())
 
         self.tags = [
             dict(

mage_ai/tests/api/operations/test_sessions.py

@@ -178,7 +178,9 @@ class SessionOperationTests(BaseApiTestCase):
     @patch('mage_ai.api.resources.SessionResource.AUTHENTICATION_MODE', 'ldap')
     @patch.object(LDAPConnection, 'authorize')
     @patch.object(LDAPConnection, 'authenticate')
-    async def test_ldap_login_with_role_mapping(self, mock_authenticate, mock_authorize):
+    async def test_ldap_login_with_role_mapping(
+        self, mock_authenticate, mock_authorize
+    ):
         mock_authenticate.return_value = (
             True,
             "Yami_Sukehiro",
@@ -188,7 +190,9 @@ class SessionOperationTests(BaseApiTestCase):
 
         Role.create_default_roles()
 
-        with patch.dict(os.environ, dict(LDAP_ROLES_MAPPING=json.dumps(dict(Admin=['Admin'])))):
+        with patch.dict(
+            os.environ, dict(LDAP_ROLES_MAPPING=json.dumps(dict(Admin=['Admin'])))
+        ):
             username = self.faker.email()
             operation = self.build_operation(
                 action=constants.CREATE,
@@ -208,3 +212,50 @@ class SessionOperationTests(BaseApiTestCase):
             user = User.query.filter(User.username == username).first()
             self.assertIsNotNone(user)
             self.assertEqual(user.roles_new[0].name, 'Admin')
+
+    @patch('mage_ai.api.resources.SessionResource.AUTHENTICATION_MODE', 'ldap')
+    @patch.object(LDAPConnection, 'authorize')
+    @patch.object(LDAPConnection, 'authenticate')
+    async def test_ldap_update_roles_on_login(self, mock_authenticate, mock_authorize):
+        mock_authenticate.return_value = (
+            True,
+            "Yami_Sukehiro",
+            dict(memberOf=['Admin']),
+        )
+        mock_authorize.return_value = True
+
+        Role.create_default_roles()
+
+        username = self.faker.email()
+
+        User.create(
+            username=username,
+            email=username,
+            roles_new=[Role.query.filter(Role.name == 'Admin').first()],
+        )
+
+        with patch.dict(
+            os.environ,
+            dict(
+                LDAP_ROLES_MAPPING=json.dumps(dict(Admin=['Editor'])),
+                UPDATE_ROLES_ON_LOGIN='1',
+            ),
+        ):
+            operation = self.build_operation(
+                action=constants.CREATE,
+                payload=dict(
+                    session=dict(
+                        email=username,
+                        password="black bull",
+                    )
+                ),
+                resource='sessions',
+                user=None,
+            )
+            await operation.execute()
+
+            mock_authenticate.assert_called_once_with(username, "black bull")
+            mock_authorize.assert_called_once_with("Yami_Sukehiro")
+            user = User.query.filter(User.username == username).first()
+            self.assertIsNotNone(user)
+            self.assertEqual(user.roles_new[0].name, 'Editor')

mage_ai/tests/api/resources/test_pipeline_resource.py

@@ -129,7 +129,7 @@ class PipelineResourceTest(BaseApiTestCase):
                     )
 
                     mock_get_async.assert_has_calls(
-                        [call.get_async(uuid) for uuid in [
+                        [call.get_async(uuid, repo_path=self.repo_path) for uuid in [
                             self.pipeline1.uuid,
                             self.pipeline2.uuid,
                             self.pipeline3.uuid,
@@ -172,7 +172,7 @@ class PipelineResourceTest(BaseApiTestCase):
                     )
 
                     mock_get_async.assert_has_calls(
-                        [call.get_async(uuid) for uuid in [
+                        [call.get_async(uuid, repo_path=self.repo_path) for uuid in [
                             self.pipeline3.uuid,
                         ]],
                         any_order=True,

mage_ai/tests/authentication/oauth/test_utils.py

@@ -10,7 +10,7 @@ from mage_ai.authentication.oauth.utils import (
     access_tokens_for_client,
     refresh_token_for_client,
 )
-from mage_ai.data_preparation.git.api import get_oauth_client_id
+from mage_ai.data_preparation.git.utils import get_oauth_client_id
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
 from mage_ai.tests.base_test import AsyncDBTestCase
 from mage_ai.tests.factory import create_user

mage_ai/tests/authentication/providers/test_okta.py

@@ -0,0 +1,43 @@
+import os
+import unittest
+from unittest.mock import patch
+
+from mage_ai.authentication.providers.okta import OktaProvider
+
+test_parameters = [
+    ('https://samples.auth0.com', 'test-client-id', 'test-client-secret'),
+    ('samples.auth0.com', 'test-client-id', 'test-client-secret'),
+    ('', 'test-client-id', 'test-client-secret'),
+    ('samples.auth0.com', '', ''),
+]
+
+
+class OktaProviderTest(unittest.TestCase):
+    def setUp(self):
+        self.authorization_endpoint = 'https://samples.auth0.com/authorize'
+        self.token_endpoint = 'https://samples.auth0.com/oauth/token'
+        self.userinfo_endpoint = 'https://samples.auth0.com/userinfo'
+
+    def test_okta_provider_initialization(self):
+        for url, id, secret in test_parameters:
+            with self.subTest():
+                with patch.dict(
+                    os.environ,
+                    dict(
+                        OKTA_DOMAIN_URL=url,
+                        OKTA_CLIENT_ID=id,
+                        OKTA_CLIENT_SECRET=secret,
+                    ),
+                ):
+                    if not all([url, id]):
+                        with self.assertRaises(ValueError):
+                            provider = OktaProvider()
+                    else:
+                        provider = OktaProvider()
+                        self.assertEqual(
+                            provider.authorization_endpoint, self.authorization_endpoint
+                        )
+                        self.assertEqual(provider.token_endpoint, self.token_endpoint)
+                        self.assertEqual(
+                            provider.userinfo_endpoint, self.userinfo_endpoint
+                        )

mage_ai/tests/data_preparation/models/block/test_global_data_product.py

@@ -72,6 +72,8 @@ class GlobalDataProductBlockTest(DBTestCase):
                 from_notebook=False,
                 logger=ANY,
                 logging_tags=ANY,
+                poll_interval=30,
+                remote_blocks=None,
                 variables=dict(mage=3),
             ),
         ])

mage_ai/tests/orchestration/db/models/test_oauth.py

@@ -17,7 +17,7 @@ class RoleTests(DBTestCase):
         Role.create_default_roles(
             entity=Entity.PROJECT,
             entity_id=test_entity_id,
-            prefix='test',
+            name_func=lambda x: f'test_{x}',
         )
         owner = Role.query.filter(Role.name == 'test_Owner').one_or_none()
 
@@ -29,12 +29,12 @@ class RoleTests(DBTestCase):
         Role.create_default_roles(
             entity=Entity.PROJECT,
             entity_id=test_entity_id,
-            prefix='test-2',
+            name_func=lambda x: f'test-2_{x}',
         )
         Role.create_default_roles(
             entity=Entity.PROJECT,
             entity_id=test_entity_id2,
-            prefix='test-2',
+            name_func=lambda x: f'test-2_{x}',
         )
         owner = Role.query.filter(Role.name == 'test-2_Owner').one_or_none()
 

mage_ai/tests/orchestration/queue/test_process_queue.py

@@ -13,6 +13,7 @@ class ProcessQueueTests(TestCase):
     def setUp(self):
         queue_config = QueueConfig.load(config=dict(concurrency=100))
         self.queue = ProcessQueue(queue_config=queue_config)
+        self.queue.start()
 
     def test_init(self):
         self.assertEqual(self.queue.size, 100)