Updated Test

superagi/resource_manager/file_manager.py

@@ -7,11 +7,14 @@
 from superagi.models.agent import Agent
 from superagi.models.agent_execution import AgentExecution
 from superagi.types.storage_types import StorageType
+
+
 class FileManager:
     def __init__(self, session: Session, agent_id: int = None, agent_execution_id: int = None):
         self.session = session
         self.agent_id = agent_id
         self.agent_execution_id = agent_execution_id
+
     def write_binary_file(self, file_name: str, data):
         if self.agent_id is not None:
             final_path = ResourceHelper.get_agent_write_resource_path(file_name,
@@ -31,21 +34,20 @@ def write_binary_file(self, file_name: str, data):
             return f"Binary {file_name} saved successfully"
         except Exception as err:
             return f"Error write_binary_file: {err}"
+
     def write_to_s3(self, file_name, final_path):
         with open(final_path, 'rb') as img:
             resource = ResourceHelper.make_written_file_resource(file_name=file_name,
                                                                  agent=Agent.get_agent_from_id(self.session,
                                                                                                self.agent_id),
                                                                  agent_execution=AgentExecution
                                                                  .get_agent_execution_from_id(self.session,
-                                                                                              self.agent_execution_id),session=self.session)
-            # if resource is not None:
-            #     self.session.add(resource)
-            #     self.session.commit()
-            #     self.session.flush()
+                                                                                              self.agent_execution_id),
+                                                                 session=self.session)
             if resource.storage_type == StorageType.S3.value:
                 s3_helper = S3Helper()
                 s3_helper.upload_file(img, path=resource.path)
+
     def write_file(self, file_name: str, content):
         if self.agent_id is not None:
             final_path = ResourceHelper.get_agent_write_resource_path(file_name,
@@ -65,6 +67,7 @@ def write_file(self, file_name: str, content):
             return f"{file_name} - File written successfully"
         except Exception as err:
             return f"Error write_file: {err}"
+
     def write_csv_file(self, file_name: str, csv_data):
         if self.agent_id is not None:
             final_path = ResourceHelper.get_agent_write_resource_path(file_name,
@@ -85,14 +88,14 @@ def write_csv_file(self, file_name: str, csv_data):
             return f"{file_name} - File written successfully"
         except Exception as err:
             return f"Error write_csv_file: {err}"
-
 
     def get_agent_resource_path(self, file_name: str):
         return ResourceHelper.get_agent_write_resource_path(file_name, agent=Agent.get_agent_from_id(self.session,
                                                                                                      self.agent_id),
                                                             agent_execution=AgentExecution
                                                             .get_agent_execution_from_id(self.session,
                                                                                          self.agent_execution_id))
+
     def read_file(self, file_name: str):
         if self.agent_id is not None:
             final_path = self.get_agent_resource_path(file_name)
@@ -106,13 +109,14 @@ def read_file(self, file_name: str):
             return content
         except Exception as err:
             return f"Error while reading file {file_name}: {err}"
+
     def get_files(self):
         """
         Gets all file names generated by the CodingTool.
         Returns:
             A list of file names.
         """
-        
+
         if self.agent_id is not None:
             final_path = self.get_agent_resource_path("")
         else:

tests/unit_tests/helper/test_resource_helper.py

@@ -1,21 +1,39 @@
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 from superagi.helper.resource_helper import ResourceHelper
 from superagi.models.agent import Agent
 from superagi.models.agent_execution import AgentExecution
+from superagi.models.resource import Resource
 
 
 def test_make_written_file_resource(mocker):
     mocker.patch('os.getcwd', return_value='/')
     mocker.patch('os.makedirs', return_value=None)
     mocker.patch('os.path.getsize', return_value=1000)
     mocker.patch('os.path.splitext', return_value=("", ".txt"))
-    mocker.patch('superagi.helper.resource_helper.get_config', side_effect=['FILE','/','/','FILE'])
+    mocker.patch('superagi.helper.resource_helper.get_config', side_effect=['FILE', '/', '/', 'FILE'])
     mock_agent = Agent(id=1, name='TestAgent')
     mock_agent_execution = AgentExecution(id=1, name='TestExecution')
+    session = MagicMock()
 
     with patch('superagi.helper.resource_helper.logger') as logger_mock:
-        result = ResourceHelper.make_written_file_resource('test.txt', mock_agent, mock_agent_execution)
+        # session.query.return_value.filter.return_value.first.return_value = None
+        session.query.return_value.filter_by.return_value.first.return_value = None
+        # Create a Resource object
+        resource = Resource(
+            name='test.txt',
+            path='/test.txt',
+            storage_type='FILE',
+            size=1000,
+            type='application/txt',
+            channel='OUTPUT',
+            agent_id=1,
+            agent_execution_id=1
+        )
+
+        # Mock the session.add() method to return the created Resource object
+        session.add.return_value = resource
+        result = ResourceHelper.make_written_file_resource('test.txt', mock_agent, mock_agent_execution, session)
 
     assert result.name == 'test.txt'
     assert result.path == '/test.txt'