mage-ai 0.9.74__py3-none-any.whl → 0.9.79__py3-none-any.whl

mage_ai/tests/data_preparation/test_variable_manager.py

@@ -1,5 +1,3 @@
-from unittest.mock import patch
-
 import pandas as pd
 from pandas.testing import assert_frame_equal
 
@@ -10,12 +8,10 @@ from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.data_preparation.variable_manager import (
     VariableManager,
     get_global_variable,
-    get_global_variables,
     set_global_variable,
 )
 from mage_ai.settings.repo import set_repo_path
 from mage_ai.tests.base_test import DBTestCase
-from mage_ai.tests.shared.mixins import ProjectPlatformMixin
 
 
 class VariableManagerTest(DBTestCase):
@@ -128,50 +124,3 @@ class VariableManagerTest(DBTestCase):
     #         get_global_variables(None, pipeline=pipeline), pipeline.variables
     #     )
     #     self.assertEqual(get_global_variables(pipeline.uuid), pipeline.variables)
-
-
-class VariableManagerProjectPlatformTests(ProjectPlatformMixin):
-    def test_get_global_variable(self):
-        with patch(
-            "mage_ai.data_preparation.variable_manager.project_platform_activated",
-            lambda: True,
-        ):
-            with patch(
-                "mage_ai.data_preparation.models.pipeline.project_platform_activated",
-                lambda: True,
-            ):
-                for settings in self.repo_paths.values():
-                    pipeline = Pipeline.create(
-                        self.faker.unique.name(),
-                        repo_path=settings["full_path"],
-                    )
-                    value = self.faker.unique.name()
-                    pipeline.variables = dict(mage=value)
-                    pipeline.save()
-
-                    self.assertEqual(get_global_variable(pipeline.uuid, "mage"), value)
-
-    def test_get_global_variables(self):
-        with patch(
-            "mage_ai.data_preparation.variable_manager.project_platform_activated",
-            lambda: True,
-        ):
-            with patch(
-                "mage_ai.data_preparation.models.pipeline.project_platform_activated",
-                lambda: True,
-            ):
-                for settings in self.repo_paths.values():
-                    pipeline = Pipeline.create(
-                        self.faker.unique.name(),
-                        repo_path=settings["full_path"],
-                    )
-                    pipeline.variables = self.faker.unique.name()
-                    pipeline.save()
-
-                    self.assertEqual(
-                        get_global_variables(None, pipeline=pipeline),
-                        pipeline.variables,
-                    )
-                    self.assertEqual(
-                        get_global_variables(pipeline.uuid), pipeline.variables
-                    )

mage_ai/tests/data_preparation/test_variable_manager_project_platform.py

@@ -0,0 +1,41 @@
+from unittest.mock import patch
+
+from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.variable_manager import (
+    get_global_variable,
+    get_global_variables,
+)
+from mage_ai.tests.shared.mixins import ProjectPlatformMixin
+
+
+@patch('mage_ai.data_preparation.models.pipeline.project_platform_activated', lambda: True)
+@patch('mage_ai.data_preparation.variable_manager.project_platform_activated', lambda: True)
+class VariableManagerProjectPlatformTests(ProjectPlatformMixin):
+    def test_get_global_variable(self):
+        for settings in self.repo_paths.values():
+            pipeline = Pipeline.create(
+                self.faker.unique.name(),
+                repo_path=settings["full_path"],
+            )
+            value = self.faker.unique.name()
+            pipeline.variables = dict(mage=value)
+            pipeline.save()
+
+            self.assertEqual(get_global_variable(pipeline.uuid, "mage"), value)
+
+    def test_get_global_variables(self):
+        for settings in self.repo_paths.values():
+            pipeline = Pipeline.create(
+                self.faker.unique.name(),
+                repo_path=settings["full_path"],
+            )
+            pipeline.variables = self.faker.unique.name()
+            pipeline.save()
+
+            self.assertEqual(
+                get_global_variables(None, pipeline=pipeline),
+                pipeline.variables,
+            )
+            self.assertEqual(
+                get_global_variables(pipeline.uuid), pipeline.variables
+            )

mage_ai/tests/io/create_table/test_postgresql.py

@@ -1,7 +1,9 @@
 import pandas as pd
+import simplejson
 
 from mage_ai.io.export_utils import infer_dtypes
 from mage_ai.io.postgres import Postgres
+from mage_ai.shared.parsers import encode_complex
 from mage_ai.tests.base_test import DBTestCase
 
 
@@ -26,3 +28,29 @@ class TestTablePostgres(DBTestCase):
         )
         self.assertEqual('CREATE TABLE Test.Test ("varchar_time" text,"datetime_time" TIMESTAMP);',
                          query)
+
+    def test_clean_array_value(self):
+        test_cases = [
+            [],
+            [123],
+            [['àabc', 'deèéf']],
+            [['08:00', '12:00'], ['15:00', '20:00']],
+            [['08:00', '12:00'], []],
+        ]
+        expected = [
+            '{}',
+            '{123}',
+            '{{"àabc", "deèéf"}}',
+            '{{"08:00", "12:00"}, {"15:00", "20:00"}}',
+            '{{"08:00", "12:00"}, {}}',
+        ]
+        for val, expected_val in zip(test_cases, expected):
+            cleaned_array_value = Postgres._clean_array_value(
+                simplejson.dumps(
+                    val,
+                    default=encode_complex,
+                    ensure_ascii=False,
+                    ignore_nan=True,
+                )
+            )
+            self.assertEqual(cleaned_array_value, expected_val)

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

@@ -3,7 +3,6 @@ from datetime import datetime, timedelta, timezone
 from unittest.mock import Mock, patch
 from uuid import uuid4
 
-from croniter import croniter
 from freezegun import freeze_time
 
 from mage_ai.data_preparation.models.block import Block
@@ -22,6 +21,7 @@ from mage_ai.orchestration.db.models.schedules import (
     PipelineSchedule,
 )
 from mage_ai.orchestration.pipeline_scheduler import configure_pipeline_run_payload
+from mage_ai.shared.croniter import croniter
 
 # from mage_ai.settings.utils import base_repo_path
 from mage_ai.shared.hash import merge_dict

mage_ai/tests/orchestration/notification/test_config.py

@@ -2,10 +2,10 @@ from mage_ai.orchestration.notification.config import NotificationConfig
 from mage_ai.tests.base_test import TestCase
 from mage_ai.tests.orchestration.notification.constants import (
     EMAIL_NOTIFICATION_CONFIG,
+    GOOGLE_CHAT_NOTIFICATION_CONFIG,
+    OPSGENIE_NOTIFICATION_CONFIG,
     SLACK_NOTIFICATION_CONFIG,
     TEAMS_NOTIFICATION_CONFIG,
-    GOOGLE_CHAT_NOTIFICATION_CONFIG,
-    OPSGENIE_NOTIFICATION_CONFIG
 )
 
 
@@ -43,7 +43,7 @@ class NotificationConfigTests(TestCase):
         self.assertIsNone(config3.email_config)
         self.assertEqual(config3.slack_config.webhook_url, 'test_webhook_url')
 
-        self.assertEqual(config4.teams_config.webhook_url, 'test_webhook_url')
+        self.assertEqual(config4.teams_config.webhook_url, ['test_webhook_url'])
 
         self.assertEqual(config5.google_chat_config.webhook_url, 'test_webhook_url')
 

mage_ai/tests/orchestration/notification/test_sender.py

@@ -112,9 +112,11 @@ class NotificationSenderTests(DBTestCase):
             f'`{pipeline_run.pipeline_schedule.name}` '
             f'at execution time `{pipeline_run.execution_date}`. Error: None'
         )
+        title = 'Failed to run Mage pipeline test_pipeline'
         mock_send_teams_message.assert_called_once_with(
             notification_config.teams_config,
             message,
+            title
         )
 
     @patch('mage_ai.orchestration.notification.sender.send_teams_message')
@@ -131,9 +133,11 @@ class NotificationSenderTests(DBTestCase):
             f'`{pipeline_run.pipeline_schedule.name}` '
             f'at execution time `{pipeline_run.execution_date}`.'
         )
+        title = 'Successfully ran Pipeline test_pipeline'
         mock_send_teams_message.assert_called_once_with(
             notification_config.teams_config,
             message,
+            title
         )
 
     @patch('mage_ai.orchestration.notification.sender.send_teams_message')
@@ -162,4 +166,4 @@ class NotificationSenderTests(DBTestCase):
             notification_config.opsgenie_config,
             message=ANY,
             description=ANY,
-        )
+        )

mage_ai/tests/orchestration/utils/test_resources.py

@@ -0,0 +1,235 @@
+import os
+import tempfile
+import unittest
+from unittest.mock import MagicMock, patch
+
+from mage_ai.orchestration.utils.resources import get_memory
+
+
+class TestGetMemory(unittest.TestCase):
+    """Test cases for the get_memory function."""
+
+    @patch('os.path.exists')
+    @patch('builtins.open', create=True)
+    def test_get_memory_cgroup_v1(self, mock_open, mock_exists):
+        """Test memory reading from cgroup v1 files."""
+        # Setup: cgroup v1 paths exist
+        def exists_side_effect(path):
+            return path in [
+                "/sys/fs/cgroup/memory/memory.usage_in_bytes",
+                "/sys/fs/cgroup/memory/memory.limit_in_bytes"
+            ]
+        
+        mock_exists.side_effect = exists_side_effect
+        
+        # Mock file reads: 1387 MB usage, 4768 MB limit (from the issue)
+        usage_bytes = 1387 * 1024 * 1024
+        limit_bytes = 4768 * 1024 * 1024
+        
+        mock_file_handles = {
+            "/sys/fs/cgroup/memory/memory.usage_in_bytes": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: str(usage_bytes)
+            ),
+            "/sys/fs/cgroup/memory/memory.limit_in_bytes": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: str(limit_bytes)
+            ),
+        }
+        
+        def open_side_effect(path, mode='r'):
+            return mock_file_handles.get(path)
+        
+        mock_open.side_effect = open_side_effect
+        
+        # Execute
+        free_memory, used_memory, total_memory = get_memory()
+        
+        # Assert
+        self.assertAlmostEqual(used_memory, 1387.0, places=1)
+        self.assertAlmostEqual(total_memory, 4768.0, places=1)
+        self.assertAlmostEqual(free_memory, 3381.0, places=1)
+
+    @patch('os.path.exists')
+    @patch('builtins.open', create=True)
+    def test_get_memory_cgroup_v2(self, mock_open, mock_exists):
+        """Test memory reading from cgroup v2 files."""
+        # Setup: cgroup v2 paths exist, v1 don't
+        def exists_side_effect(path):
+            return path in [
+                "/sys/fs/cgroup/memory.current",
+                "/sys/fs/cgroup/memory.max"
+            ]
+        
+        mock_exists.side_effect = exists_side_effect
+        
+        # Mock file reads
+        usage_bytes = 1500 * 1024 * 1024
+        limit_bytes = 5000 * 1024 * 1024
+        
+        # Create mock read() return value with strip() method
+        usage_read_value = MagicMock()
+        usage_read_value.strip.return_value = str(usage_bytes)
+        
+        limit_read_value = MagicMock()
+        limit_read_value.strip.return_value = str(limit_bytes)
+        
+        mock_file_handles = {
+            "/sys/fs/cgroup/memory.current": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: usage_read_value
+            ),
+            "/sys/fs/cgroup/memory.max": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: limit_read_value
+            ),
+        }
+        
+        def open_side_effect(path, mode='r'):
+            return mock_file_handles.get(path)
+        
+        mock_open.side_effect = open_side_effect
+        
+        # Execute
+        free_memory, used_memory, total_memory = get_memory()
+        
+        # Assert
+        self.assertAlmostEqual(used_memory, 1500.0, places=1)
+        self.assertAlmostEqual(total_memory, 5000.0, places=1)
+        self.assertAlmostEqual(free_memory, 3500.0, places=1)
+
+    @patch('os.path.exists')
+    @patch('builtins.open', create=True)
+    def test_get_memory_cgroup_v2_max_limit(self, mock_open, mock_exists):
+        """Test memory reading from cgroup v2 with 'max' limit (no limit)."""
+        # Setup: cgroup v2 with max limit
+        def exists_side_effect(path):
+            return path in [
+                "/sys/fs/cgroup/memory.current",
+                "/sys/fs/cgroup/memory.max"
+            ]
+        
+        mock_exists.side_effect = exists_side_effect
+        
+        usage_bytes = 1500 * 1024 * 1024
+        
+        # Create mock read() return value with strip() method
+        usage_read_value = MagicMock()
+        usage_read_value.strip.return_value = str(usage_bytes)
+        
+        limit_read_value = MagicMock()
+        limit_read_value.strip.return_value = "max"
+        
+        mock_file_handles = {
+            "/sys/fs/cgroup/memory.current": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: usage_read_value
+            ),
+            "/sys/fs/cgroup/memory.max": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: limit_read_value
+            ),
+        }
+        
+        def open_side_effect(path, mode='r'):
+            return mock_file_handles.get(path)
+        
+        mock_open.side_effect = open_side_effect
+        
+        # With 'max' limit, should fall back to free command
+        # We'll mock subprocess as well
+        with patch('subprocess.check_output') as mock_subprocess:
+            mock_subprocess.return_value = b"Mem:     8000     4000     4000\nSwap:       0        0        0\nTotal:   8000     4000     4000"
+            
+            free_memory, used_memory, total_memory = get_memory()
+            
+            # Should use free command fallback
+            self.assertEqual(total_memory, 8000.0)
+            self.assertEqual(used_memory, 4000.0)
+            self.assertEqual(free_memory, 4000.0)
+
+    @patch('os.path.exists')
+    @patch('subprocess.check_output')
+    def test_get_memory_fallback_to_free(self, mock_subprocess, mock_exists):
+        """Test fallback to 'free' command when cgroup files don't exist."""
+        # Setup: no cgroup files exist
+        mock_exists.return_value = False
+        
+        # Mock free command output
+        mock_subprocess.return_value = b"Mem:     4570     1921     2533\nSwap:       0        0        0\nTotal:   4570     1921     2533"
+        
+        # Execute
+        free_memory, used_memory, total_memory = get_memory()
+        
+        # Assert - values from the issue's screenshot
+        self.assertEqual(total_memory, 4570.0)
+        self.assertEqual(used_memory, 1921.0)
+        self.assertEqual(free_memory, 2533.0)
+
+    @patch('os.name', 'nt')
+    def test_get_memory_windows(self):
+        """Test that Windows returns None values."""
+        free_memory, used_memory, total_memory = get_memory()
+        
+        self.assertIsNone(free_memory)
+        self.assertIsNone(used_memory)
+        self.assertIsNone(total_memory)
+
+    @patch('os.path.exists')
+    @patch('builtins.open', create=True)
+    def test_get_memory_cgroup_v1_no_limit(self, mock_open, mock_exists):
+        """Test memory reading with very large limit value (no limit set)."""
+        # Setup: cgroup v1 paths exist
+        def exists_side_effect(path):
+            return path in [
+                "/sys/fs/cgroup/memory/memory.usage_in_bytes",
+                "/sys/fs/cgroup/memory/memory.limit_in_bytes"
+            ]
+        
+        mock_exists.side_effect = exists_side_effect
+        
+        usage_bytes = 1500 * 1024 * 1024
+        # Very large limit indicates no limit
+        limit_bytes = 9223372036854771712
+        
+        # Create mock read() return value with strip() method
+        usage_read_value = MagicMock()
+        usage_read_value.strip.return_value = str(usage_bytes)
+        
+        limit_read_value = MagicMock()
+        limit_read_value.strip.return_value = str(limit_bytes)
+        
+        mock_file_handles = {
+            "/sys/fs/cgroup/memory/memory.usage_in_bytes": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: usage_read_value
+            ),
+            "/sys/fs/cgroup/memory/memory.limit_in_bytes": MagicMock(
+                __enter__=lambda self: self,
+                __exit__=lambda self, *args: None,
+                read=lambda: limit_read_value
+            ),
+        }
+        
+        def open_side_effect(path, mode='r'):
+            return mock_file_handles.get(path)
+        
+        mock_open.side_effect = open_side_effect
+        
+        # With no real limit, should fall back to free command
+        with patch('subprocess.check_output') as mock_subprocess:
+            mock_subprocess.return_value = b"Mem:     8000     4000     4000\nSwap:       0        0        0\nTotal:   8000     4000     4000"
+            
+            free_memory, used_memory, total_memory = get_memory()
+            
+            # Should use free command fallback
+            self.assertEqual(total_memory, 8000.0)
+            self.assertEqual(used_memory, 4000.0)
+            self.assertEqual(free_memory, 4000.0)