mage-ai 0.9.65__py3-none-any.whl → 0.9.67__py3-none-any.whl

mage_ai/tests/api/operations/test_sessions.py

@@ -1,4 +1,5 @@
 import json
+import os
 from datetime import datetime
 from unittest.mock import patch
 
@@ -20,15 +21,20 @@ class SessionOperationTests(BaseApiTestCase):
 
         operation = self.build_operation(
             action=constants.CREATE,
-            payload=dict(session=dict(
-                email=user.email,
-                password=password,
-            )),
+            payload=dict(
+                session=dict(
+                    email=user.email,
+                    password=password,
+                )
+            ),
             resource='sessions',
+            user=None,
         )
         response = await operation.execute()
 
-        access_token = Oauth2AccessToken.query.filter(Oauth2AccessToken.user_id == user.id).first()
+        access_token = Oauth2AccessToken.query.filter(
+            Oauth2AccessToken.user_id == user.id
+        ).first()
 
         self.assertEqual(
             response['session']['token'],
@@ -41,22 +47,32 @@ class SessionOperationTests(BaseApiTestCase):
         user = create_user(password=password)
 
         with patch('mage_ai.api.policies.BasePolicy.DISABLE_NOTEBOOK_EDIT_ACCESS', 1):
-            with patch('mage_ai.api.policies.BasePolicy.REQUIRE_USER_AUTHENTICATION', 1):
-                with patch('mage_ai.api.policies.BasePolicy.REQUIRE_USER_PERMISSIONS', 0):
+            with patch(
+                'mage_ai.api.policies.BasePolicy.REQUIRE_USER_AUTHENTICATION', 1
+            ):
+                with patch(
+                    'mage_ai.api.policies.BasePolicy.REQUIRE_USER_PERMISSIONS', 0
+                ):
                     operation = self.build_operation(
                         action=constants.CREATE,
-                        payload=dict(session=dict(
-                            email=user.email,
-                            password=password,
-                        )),
+                        payload=dict(
+                            session=dict(
+                                email=user.email,
+                                password=password,
+                            )
+                        ),
                         resource='sessions',
                         user=None,
                     )
                     response = await operation.execute()
 
-                    access_token = Oauth2AccessToken.query.filter(
-                        Oauth2AccessToken.user_id == user.id
-                    ).order_by(Oauth2AccessToken.created_at.desc()).first()
+                    access_token = (
+                        Oauth2AccessToken.query.filter(
+                            Oauth2AccessToken.user_id == user.id
+                        )
+                        .order_by(Oauth2AccessToken.created_at.desc())
+                        .first()
+                    )
 
                     self.assertEqual(
                         response['session']['token'],
@@ -69,11 +85,14 @@ class SessionOperationTests(BaseApiTestCase):
 
         operation = self.build_operation(
             action=constants.CREATE,
-            payload=dict(session=dict(
-                email=user.email,
-                password='not password',
-            )),
+            payload=dict(
+                session=dict(
+                    email=user.email,
+                    password='not password',
+                )
+            ),
             resource='sessions',
+            user=None,
         )
         response = await operation.execute()
 
@@ -89,11 +108,14 @@ class SessionOperationTests(BaseApiTestCase):
         username = "novachrono"
         operation = self.build_operation(
             action=constants.CREATE,
-            payload=dict(session=dict(
-                email=username,
-                password="Wizard King",
-            )),
+            payload=dict(
+                session=dict(
+                    email=username,
+                    password="Wizard King",
+                )
+            ),
             resource='sessions',
+            user=None,
         )
         await operation.execute()
 
@@ -111,11 +133,14 @@ class SessionOperationTests(BaseApiTestCase):
         username = "licht"
         operation = self.build_operation(
             action=constants.CREATE,
-            payload=dict(session=dict(
-                email=username,
-                password="Golden Dawn",
-            )),
+            payload=dict(
+                session=dict(
+                    email=username,
+                    password="Golden Dawn",
+                )
+            ),
             resource='sessions',
+            user=None,
         )
         response = await operation.execute()
 
@@ -134,11 +159,14 @@ class SessionOperationTests(BaseApiTestCase):
         username = "yami"
         operation = self.build_operation(
             action=constants.CREATE,
-            payload=dict(session=dict(
-                email=username,
-                password="black bull",
-            )),
+            payload=dict(
+                session=dict(
+                    email=username,
+                    password="black bull",
+                )
+            ),
             resource='sessions',
+            user=None,
         )
         response = await operation.execute()
 
@@ -148,28 +176,35 @@ class SessionOperationTests(BaseApiTestCase):
         self.assertIsNotNone(response['error'])
 
     @patch('mage_ai.api.resources.SessionResource.AUTHENTICATION_MODE', 'ldap')
-    @patch('mage_ai.authentication.ldap.LDAP_ROLES_MAPPING', json.dumps(dict(Admin=['Admin'])))
     @patch.object(LDAPConnection, 'authorize')
     @patch.object(LDAPConnection, 'authenticate')
     async def test_ldap_login_with_role_mapping(self, mock_authenticate, mock_authorize):
-        mock_authenticate.return_value = (True, "Yami_Sukehiro", dict(memberOf=['Admin']))
+        mock_authenticate.return_value = (
+            True,
+            "Yami_Sukehiro",
+            dict(memberOf=['Admin']),
+        )
         mock_authorize.return_value = True
 
         Role.create_default_roles()
 
-        username = self.faker.email()
-        operation = self.build_operation(
-            action=constants.CREATE,
-            payload=dict(session=dict(
-                email=username,
-                password="black bull",
-            )),
-            resource='sessions',
-        )
-        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, '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,
+                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, 'Admin')

mage_ai/tests/authentication/oauth/test_utils.py

@@ -1,14 +1,19 @@
 import asyncio
 from datetime import datetime, timedelta
 from unittest.mock import MagicMock, patch
+from uuid import uuid4
 
 from freezegun import freeze_time
 
 from mage_ai.authentication.oauth.constants import ProviderName
-from mage_ai.authentication.oauth.utils import refresh_token_for_client
+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.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
 from mage_ai.tests.base_test import AsyncDBTestCase
+from mage_ai.tests.factory import create_user
 
 
 class OauthUtilsTest(AsyncDBTestCase):
@@ -17,11 +22,54 @@ class OauthUtilsTest(AsyncDBTestCase):
         super().setUpClass()
         self.provider = ProviderName.GITLAB
         self.client_id = get_oauth_client_id(self.provider)
+
+    def setUp(self):
         self.oauth_application = Oauth2Application.create(
             client_id=self.client_id,
             client_type=Oauth2Application.ClientType.PRIVATE,
             name=self.provider,
         )
+        self.user = create_user()
+
+    def tearDown(self):
+        Oauth2Application.query.delete()
+        Oauth2AccessToken.query.delete()
+
+        super().tearDown()
+
+    def test_access_tokens_for_client(self):
+        token = Oauth2AccessToken.create(
+            expires=datetime.utcnow() + timedelta(days=1),
+            token=uuid4().hex,
+            oauth2_application_id=self.oauth_application.id,
+        )
+
+        tokens = access_tokens_for_client(self.client_id)
+
+        self.assertEqual(len(tokens), 1)
+        self.assertEqual(tokens[0].token, token.token)
+
+    def test_access_tokens_for_client_with_user(self):
+        Oauth2AccessToken.create(
+            expires=datetime.utcnow() + timedelta(days=1),
+            token=uuid4().hex,
+            oauth2_application_id=self.oauth_application.id,
+        )
+
+        self.assertEqual(
+            len(access_tokens_for_client(self.client_id, user=self.user)), 0
+        )
+        self.assertEqual(len(access_tokens_for_client(self.client_id)), 1)
+
+        token = Oauth2AccessToken.create(
+            expires=datetime.utcnow() + timedelta(days=1),
+            token=uuid4().hex,
+            oauth2_application_id=self.oauth_application.id,
+            user_id=self.user.id,
+        )
+        tokens = access_tokens_for_client(self.client_id, user=self.user)
+        self.assertEqual(len(tokens), 1)
+        self.assertEqual(tokens[0].token, token.token)
 
     @freeze_time('2024-02-06 04:00:00')
     @patch(
@@ -43,11 +91,13 @@ class OauthUtilsTest(AsyncDBTestCase):
         provider_instance.provider = self.provider
 
         future = asyncio.Future()
-        future.set_result({
-            'access_token': 'new',
-            'refresh_token': 'new_refresh',
-            'expires_in': 3600,
-        })
+        future.set_result(
+            {
+                'access_token': 'new',
+                'refresh_token': 'new_refresh',
+                'expires_in': 3600,
+            }
+        )
         provider_instance.get_refresh_token_response.return_value = future
         new_token = await refresh_token_for_client(self.client_id, provider_instance)
 

mage_ai/tests/authentication/providers/test_active_directory.py

@@ -1,4 +1,5 @@
 import json
+import os
 import unittest
 from unittest.mock import AsyncMock, Mock, patch
 
@@ -6,18 +7,6 @@ from mage_ai.authentication.providers.active_directory import ADProvider
 
 
 class ADProviderTest(unittest.IsolatedAsyncioTestCase):
-    @patch(
-        'mage_ai.authentication.providers.active_directory.ACTIVE_DIRECTORY_DIRECTORY_ID',
-        'test_directory_id',
-    )
-    @patch(
-        'mage_ai.authentication.providers.active_directory.ACTIVE_DIRECTORY_CLIENT_ID',
-        'test_client_id',
-    )
-    @patch(
-        'mage_ai.authentication.providers.active_directory.ACTIVE_DIRECTORY_CLIENT_SECRET',
-        'test_client_secret',
-    )
     @patch('aiohttp.ClientSession')
     async def test_map_ad_roles_to_mage_roles(self, mock_client_session):
         roles_mapping = json.dumps(
@@ -27,9 +16,14 @@ class ADProviderTest(unittest.IsolatedAsyncioTestCase):
             }
         )
 
-        with patch(
-            'mage_ai.authentication.providers.active_directory.ACTIVE_DIRECTORY_ROLES_MAPPING',
-            roles_mapping,
+        with patch.dict(
+            os.environ,
+            dict(
+                ACTIVE_DIRECTORY_DIRECTORY_ID='test_directory_id',
+                ACTIVE_DIRECTORY_CLIENT_ID='test_client_id',
+                ACTIVE_DIRECTORY_CLIENT_SECRET='test_client_secret',
+                ACTIVE_DIRECTORY_ROLES_MAPPING=roles_mapping,
+            ),
         ):
 
             def side_effect(*args, **kwargs):

mage_ai/tests/data_preparation/models/test_block.py

@@ -7,14 +7,17 @@ from faker import Faker
 from pandas.testing import assert_frame_equal
 
 # from mage_ai.data_cleaner.column_types.constants import ColumnType
-from mage_ai.data_preparation.models.block import Block, BlockType
+from mage_ai.data_preparation.models.block import Block, BlockType, CallbackBlock
 from mage_ai.data_preparation.models.block.errors import HasDownstreamDependencies
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.data_preparation.variable_manager import VariableManager
 from mage_ai.shared.path_fixer import add_root_repo_path_to_relative_path
 from mage_ai.tests.base_test import DBTestCase
-from mage_ai.tests.factory import create_integration_pipeline_with_blocks
+from mage_ai.tests.factory import (
+    create_integration_pipeline_with_blocks,
+    create_pipeline,
+)
 from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 
@@ -763,3 +766,37 @@ class BlockProjectPlatformTests(ProjectPlatformMixin):
                 )
 
                 self.assertEqual(block.file_path, add_root_repo_path_to_relative_path(path))
+
+
+class CallbackBlockTest(DBTestCase):
+    def setUp(self):
+        self.pipeline = create_pipeline('callback_pipeline', self.repo_path)
+
+    def tearDown(self):
+        self.pipeline.delete()
+
+    def test_create_global_vars_from_parent_block(self):
+        parent_block = Block.create(
+            'test_data_loader',
+            'data_loader',
+            self.repo_path,
+            pipeline=self.pipeline,
+        )
+        callback_block = CallbackBlock.create(parent_block.name)
+        self.pipeline.add_block(callback_block)
+        parent_block = parent_block.update(
+            dict(callback_blocks=[callback_block.uuid])
+        )
+        parent_block.global_vars = dict(
+            configuration=dict(table_name='load_data_table')
+        )
+
+        global_vars = dict(
+            random_var=1,
+        )
+        new_vars = callback_block._create_global_vars(
+            global_vars,
+            parent_block=parent_block,
+        )
+
+        self.assertIsNotNone(new_vars.get('configuration'))

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

@@ -1,5 +1,5 @@
 import os
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from unittest.mock import Mock, patch
 from uuid import uuid4
 
@@ -1677,6 +1677,38 @@ class PipelineRunTests(DBTestCase):
             set([pipeline_run2.id, pipeline_run4.id]),
         )
 
+    @freeze_time('2023-10-11 12:13:14')
+    def test_get_variables_intervals_for_cron(self):
+        pipeline_schedule = PipelineSchedule.create(
+            name=self.faker.unique.name(),
+            pipeline_uuid=self.pipeline.uuid,
+            schedule_interval='*/10 * * * *',
+            schedule_type=ScheduleType.TIME,
+        )
+
+        pipeline_run = PipelineRun.create(
+            pipeline_uuid=self.pipeline.uuid,
+            pipeline_schedule_id=pipeline_schedule.id,
+            create_block_runs=False,
+            execution_date=pipeline_schedule.current_execution_date(),
+        )
+
+        variables = pipeline_run.get_variables(pipeline_uuid=pipeline_run.pipeline_uuid)
+
+        self.assertEqual(
+            datetime(2023, 10, 11, 12, 20, 0) - variables['interval_end_datetime'],
+            timedelta(0),
+        )
+        self.assertEqual(600, variables['interval_seconds'])
+        self.assertEqual(
+            datetime(2023, 10, 11, 12, 10, 0) - variables['interval_start_datetime'],
+            timedelta(0),
+        )
+        self.assertEqual(
+            datetime(2023, 10, 11, 12, 0, 0) - variables['interval_start_datetime_previous'],
+            timedelta(0),
+        )
+
 
 class BlockRunTests(DBTestCase):
     @classmethod

{mage_ai-0.9.65.dist-info → mage_ai-0.9.67.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.65
+Version: 0.9.67
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -58,6 +58,7 @@ Requires-Dist: thefuzz[speedup] ==0.19.0
 Requires-Dist: tornado ==6.3.3
 Requires-Dist: typer[all] ==0.9.0
 Requires-Dist: typing-extensions ==4.5.0
+Requires-Dist: watchdog ==4.0.0
 Provides-Extra: ai
 Requires-Dist: astor >=0.8.1 ; extra == 'ai'
 Requires-Dist: langchain >=0.0.222 ; extra == 'ai'