How to Automate Unit Tests and Documentation with AI Agents

1import os
2os.environ["TEAM_API_KEY"] = "TEAM_API_KEY_HERE"
3from aixplain.factories.agent_factory import AgentFactory
4from aixplain.factories.team_agent_factory import TeamAgentFactory
Copy

1import os
2def get_files_content(code_directory):
3    python_files = []
4    for r, d, f in os.walk(code_directory):
5        for file in f:
6            if file.endswith(".py"):
7                python_files.append(os.path.join(r, file))
8                
9    output = []
10    for file in python_files:
11        with open(file, "r", encoding="utf8") as f:
12            content = f.read()
13        output.append((file, content))
14    return output
Copy

1import os
2def get_files_content(code_directory):
3    python_files = []
4    for r, d, f in os.walk(code_directory):
5        for file in f:
6            if file.endswith(".py"):
7                python_files.append(os.path.join(r, file))
8                
9    output = []
10    for file in python_files:
11        with open(file, "r", encoding="utf8") as f:
12            content = f.read()
13        output.append((file, content))
14    return output
Copy

1# team agent
2team_agent = TeamAgentFactory.create(name="DocumentTeamAgent", agents=[test_agent, doc_agent], llm_id=llm_id)
Copy

1
2# get file content
3file_content = get_files_content("src/")
4if not file_content:
5    print("Folder must contain .py files")
6    return
7
8with open("prompt.txt", "r") as f:
9    prompt = f.read().strip()
10for file_path, content in file_content:
11    prompt += f"\nFile path: {file_path}\nContent: {content}"
12    prompt += "OUTPUT:\n"
13
Copy

1# run team agent
2response = team_agent.run(prompt)
Copy

1# parse the response
2if not os.path.exists("tests"):
3  os.mkdir("tests")
4with open(f"tests/test.md", "w") as f:
5  f.write(response["data"]["output"])
Copy

1# 1. File path: src/utils/config.py
2import pytest
3import os
4from unittest.mock import patch
5
6# Mocking environment variables
7@patch.dict(os.environ, {
8    "BACKEND_URL": "https://mock-backend-url.com",
9    "MODELS_RUN_URL": "https://mock-models-run-url.com",
10    "TEAM_API_KEY": "mock_team_api_key",
11    "AIXPLAIN_API_KEY": "mock_aixplain_api_key",
12    "PIPELINE_API_KEY": "mock_pipeline_api_key",
13    "MODEL_API_KEY": "mock_model_api_key",
14    "LOG_LEVEL": "DEBUG",
15    "HF_TOKEN": "mock_hf_token"
16})
17def test_config_values():
18    import src.utils.config as config
19    assert config.BACKEND_URL == "https://mock-backend-url.com"
20    assert config.MODELS_RUN_URL == "https://mock-models-run-url.com"
21    assert config.TEAM_API_KEY == "mock_team_api_key"
22    assert config.AIXPLAIN_API_KEY == "mock_aixplain_api_key"
23    assert config.PIPELINE_API_KEY == "mock_pipeline_api_key"
24    assert config.MODEL_API_KEY == "mock_model_api_key"
25    assert config.LOG_LEVEL == "DEBUG"
26    assert config.HF_TOKEN == "mock_hf_token"
Copy

1import pytest
2from unittest.mock import patch, Mock
3from requests.models import Response
4from src.utils.request_utils import _request_with_retry
5
6@patch('src.utils.request_utils.requests.Session')
7def test_request_with_retry(mock_session):
8    # Create a mock response object
9    mock_response = Mock(spec=Response)
10    mock_response.status_code = 200
11    
12    # Configure the mock session to return the mock response
13    mock_session_instance = mock_session.return_value
14    mock_session_instance.request.return_value = mock_response
Copy

1# Call the function
2response = _request_with_retry('GET', 'https://example.com')
3
4# Assert that the request was made with the correct parameters
5            mock_session_instance.request.assert_called_once_with(method='GET', url='https://example.com')
6
7# Assert that the response is as expected
8assert response == mock_response
Copy

1import pytest
2from unittest.mock import patch, mock_open, MagicMock
3from src.utils.file_utils import save_file, download_data, upload_data, s3_to_csv
4
5@patch('src.utils.file_utils._request_with_retry')
6@patch('builtins.open', new_callable=mock_open)
7@patch('os.getcwd', return_value='/mocked/path')
8@patch('os.path.basename', return_value='mocked_file.csv')
9@patch('src.utils.file_utils.config')
10def test_save_file(mock_config, mock_basename, mock_getcwd, mock_open, mock_request):
11    mock_response = MagicMock()
12    mock_response.content = b'test content'
13    mock_request.return_value = mock_response
14    
15    download_url = 'http://example.com/file.csv'
16    download_file_path = save_file(download_url)
17    
18    mock_request.assert_called_once_with('get', download_url)
19    mock_open.assert_called_once_with('/mocked/path/aiXplain/mocked_file.csv', 'wb')
20    assert download_file_path == '/mocked/path/aiXplain/mocked_file.csv'
21
22@patch('requests.get')
23@patch('builtins.open', new_callable=mock_open)
24def test_download_data(mock_open, mock_get):
25    mock_response = MagicMock()
26    mock_response.iter_content = lambda chunk_size: [b'test content']
27    mock_response.raise_for_status = lambda: None
28    mock_get.return_value = mock_response
29    
30    url_link = 'http://example.com/file.csv'
31    local_filename = download_data(url_link)
32    
33    mock_get.assert_called_once_with(url_link, stream=True)
34    mock_open.assert_called_once_with('file.csv', 'wb')
35    assert local_filename == 'file.csv'
36
37@patch('src.utils.file_utils._request_with_retry')
38@patch('builtins.open', new_callable=mock_open, read_data=b'test content')
39@patch('os.path.basename', return_value='mocked_file.csv')
40@patch('src.utils.file_utils.config')
41def test_upload_data(mock_config, mock_basename, mock_open, mock_request):
42    mock_response = MagicMock()
43    mock_response.json.return_value = {
44        'key': 'mocked_key',
45        'uploadUrl': 'http://example.com/upload'
46    }
47    mock_request.return_value = mock_response
48    
49    mock_config.BACKEND_URL = 'http://backend.example.com'
50    mock_config.AIXPLAIN_API_KEY = 'mocked_api_key'
51    
52    file_name = 'mocked_file.csv'
53    s3_link = upload_data(file_name)
54    
55    mock_request.assert_any_call('post', 'http://backend.example.com/sdk/file/upload/temp-url', headers={'x-aixplain-key': 'mocked_api_key'}, data={'contentType': 'text/csv', 'originalName': 'mocked_file.csv'})
56    mock_request.assert_any_call('put', 'http://example.com/upload', headers={'Content-Type': 'text/csv'}, data=b'test content')
57    assert s3_link == 's3://example.com/mocked_key'
58
59@patch('boto3.client')
60@patch('os.getenv', side_effect=lambda key: 'mocked_value')
61def test_s3_to_csv(mock_getenv, mock_boto3_client):
62    mock_s3 = MagicMock()
63    mock_boto3_client.return_value = mock_s3
64    mock_s3.list_objects_v2.return_value = {'Contents': [{'Key': 'mocked_file.csv'}]}
65    
66    s3_url = 's3://mocked-bucket/mocked_file.csv'
67    aws_credentials = {'AWS_ACCESS_KEY_ID': 'mocked_id', 'AWS_SECRET_ACCESS_KEY': 'mocked_secret'}
68    csv_path = s3_to_csv(s3_url, aws_credentials)
69    
70    mock_boto3_client.assert_called_once_with('s3', aws_access_key_id='mocked_id', aws_secret_access_key='mocked_secret')
71    mock_s3.list_objects_v2.assert_called_once_with(Bucket='mocked-bucket', Prefix='mocked_file.csv')
72    assert csv_path == 'mocked_file.csv'
Copy

1# Project Documentation
2
3## Overview
4This project is designed to handle various utility functions related to configuration, HTTP requests, and file operations, particularly focusing on interactions with web services and AWS S3. Below is a summary of each file's purpose and functionality within the project.
5
6### File: `src/utils/config.py`
7This file is responsible for managing configuration settings for the project. It utilizes the `os` module to fetch environment variables, which include URLs and API keys necessary for the application's operation. This setup allows for flexible configuration management, adapting to different environments by simply changing environment variables.
8
9### File: `src/utils/request_utils.py`
10This file provides a utility function named `_request_with_retry`. This function is a wrapper around the `requests` library, designed to perform HTTP requests with built-in retry logic. It uses a session configured with retry settings to handle transient errors, ensuring more robust and reliable HTTP communication.
11
12### File: `src/utils/file_utils.py`
13This file contains several utility functions aimed at facilitating file operations. Key functionalities include:
14- Downloading files from a specified URL.
15- Uploading data to AWS S3 using pre-signed URLs, which allows secure and temporary access to S3 resources.
16- Converting S3 URLs to CSV files, enabling easy data manipulation and analysis.
17
18These utilities handle various aspects of file management, including file paths, HTTP requests, and interactions with AWS S3, making it easier to integrate file operations into the project workflow.
Copy