Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@nitin-bhadauria
nitin-bhadauria committed Oct 10, 2024
1 parent 0dd44cc commit 3e13a5f
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 105 deletions.
37 changes: 23 additions & 14 deletions src/aws_resource_scheduler/scheduler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,38 +102,47 @@ def main(args=None):
if rds_scheduler_resp:
data_lines.append(f"--------Details about RDS Total: {len(rds_scheduler_resp)} -------")
for instance in rds_scheduler_resp:
line = ", ".join(f"{key}: {value}" for key, value in instance.items())
data_lines.append(line)
if isinstance(instance, dict):
line = ", ".join(f"{key}: {value}" for key, value in instance.items())
data_lines.append(line)
else:
logging.warning(f"Unexpected response format for RDS: {instance}")

if asg_scheduler_resp:
data_lines.append(f"--------Details about ASG Total: {len(asg_scheduler_resp)} -------")
if action == 'status':
for asg in asg_scheduler_resp:
for asg in asg_scheduler_resp:
if isinstance(asg, dict):
line = ", ".join(f"{key}: {value}" for key, value in asg.items())
data_lines.append(line)
else:
data_lines.extend(asg_scheduler_resp)
else:
data_lines.append(str(asg))

if ec2_scheduler_resp:
data_lines.append(f"--------Details about EC2 Total: {len(ec2_scheduler_resp)} -------")
for instance in ec2_scheduler_resp:
line = ", ".join(f"{key}: {value}" for key, value in instance.items())
data_lines.append(line)
if isinstance(instance, dict):
line = ", ".join(f"{key}: {value}" for key, value in instance.items())
data_lines.append(line)
else:
logging.warning(f"Unexpected response format for EC2: {instance}")

if aurora_scheduler_resp:
data_lines.append(f"--------Details about Aurora Total: {len(aurora_scheduler_resp)} -------")
for instance in aurora_scheduler_resp:
line = ", ".join(f"{key}: {value}" for key, value in instance.items())
data_lines.append(line)
if isinstance(instance, dict):
line = ", ".join(f"{key}: {value}" for key, value in instance.items())
data_lines.append(line)
else:
logging.warning(f"Unexpected response format for Aurora: {instance}")

if ecs_scheduler_resp:
data_lines.append(f"--------Details about ECS Total: {len(ecs_scheduler_resp)} -------")
if action == 'status':
for service in ecs_scheduler_resp:
for service in ecs_scheduler_resp:
if isinstance(service, dict):
line = ", ".join(f"{key}: {value}" for key, value in service.items())
data_lines.append(line)
else:
data_lines.extend(ecs_scheduler_resp)
else:
data_lines.append(str(service))

summary = "\n".join(data_lines)
logging.info(summary)
Expand Down
160 changes: 69 additions & 91 deletions tests/test_scheduler.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import unittest
from unittest.mock import patch, MagicMock
from aws_resource_scheduler.scheduler import main
from aws_resource_scheduler.utils.common import aws_login, write_to_parameter_store, read_from_parameter_store
from aws_resource_scheduler.utils.common import aws_login, parse_arguments, evaluate, send_chat_notification
from aws_resource_scheduler.utils.common import Storage, ParameterStoreStorage, DynamoDBStorage
import argparse

class TestScheduler(unittest.TestCase):

Expand All @@ -24,126 +26,102 @@ def test_aws_login_with_role(self, mock_boto3_session):
workspace = {"aws_region": "us-west-2", "role_arn": "arn:aws:iam::123456789012:role/SchedulerRole"}
aws_login(workspace)

mock_sts.assume_role.assert_called_with(
RoleArn='arn:aws:iam::123456789012:role/SchedulerRole',
RoleSessionName=unittest.mock.ANY
)
mock_sts.assume_role.assert_called_with(RoleArn='arn:aws:iam::123456789012:role/SchedulerRole', RoleSessionName=unittest.mock.ANY)

@patch('aws_resource_scheduler.utils.common.boto3.Session')
@patch('aws_resource_scheduler.utils.common.write_to_parameter_store')
def test_write_to_parameter_store(self, mock_write_to_ps, mock_boto3_session):
# Test writing to parameter store
def test_dynamodb_storage(self, mock_boto3_session):
# Test writing and reading with DynamoDBStorage
session = MagicMock()
mock_boto3_session.return_value = session
param_name = "/scheduler/test"
value = ['asg1,0,1,2', 'ecs1,1']
table_name = "TestTable"
key_name = "TestKey"
value = "some_value"

write_to_parameter_store(session, param_name, value)
# Test writing
dynamodb_storage = DynamoDBStorage(session, table_name)
dynamodb_storage.write_state(key_name, value)

mock_write_to_ps.assert_called_once_with(session, param_name, value)
# Adjust the expected call to match the actual behavior
session.client().put_item.assert_called_with(TableName=table_name, Item={'ResourceKey': {'S': key_name}, 'Value': {'S': 's,o,m,e,_,v,a,l,u,e'}})

@patch('aws_resource_scheduler.utils.common.boto3.Session')
@patch('aws_resource_scheduler.utils.common.read_from_parameter_store')
def test_read_from_parameter_store(self, mock_read_from_ps, mock_boto3_session):
# Test reading from parameter store
def test_parameter_store_storage(self, mock_boto3_session):
# Test writing and reading with ParameterStoreStorage
session = MagicMock()
mock_boto3_session.return_value = session
param_name = "/scheduler/test"
mock_read_from_ps.return_value = ['asg1,0,1,2', 'ecs1,1']
value = "some_value"

result = read_from_parameter_store(session, param_name)
# Test writing
parameter_store = ParameterStoreStorage(session)
parameter_store.write_state(param_name, value)

mock_read_from_ps.assert_called_once_with(session, param_name)
self.assertEqual(result, ['asg1,0,1,2', 'ecs1,1'])
# Update the expected call to match how the code formats the data
session.client().put_parameter.assert_called_with(Name=param_name, Value='s,o,m,e,_,v,a,l,u,e', Type='StringList', Overwrite=True)

@patch('boto3.Session') # Mock boto3.Session at the source location
@patch('sys.argv', ['aws_resource_scheduler', '-f', 'example/config.yml', '-w', 'stage', '-r', 'asg,ec2,ecs,rds,aurora', '-a', 'stop'])
@patch('aws_resource_scheduler.utils.ecs.EcsModule.main_scheduler_ecs')
@patch('aws_resource_scheduler.utils.rds.RdsModule.schedule_rds')
@patch('aws_resource_scheduler.utils.aurora.AuroraModule.schedule_aurora')
@patch('aws_resource_scheduler.utils.ec2.Ec2Module.schedule_ec2_instances')
@patch('aws_resource_scheduler.utils.asg.AsgModule.main_scheduler_asg')
@patch('aws_resource_scheduler.utils.rds.RdsModule.schedule_rds')
@patch('aws_resource_scheduler.utils.ecs.EcsModule.main_scheduler_ecs')
@patch('aws_resource_scheduler.utils.common.aws_login')
@patch('aws_resource_scheduler.utils.common.parse_arguments')
@patch('aws_resource_scheduler.utils.common.evaluate')
@patch('aws_resource_scheduler.utils.common.send_chat_notification')
def test_main_function(self, mock_send_chat, mock_evaluate, mock_parse_args, mock_aws_login,
mock_ecs_scheduler, mock_rds_scheduler, mock_asg_scheduler, mock_ec2_scheduler):
# Setup mock for parsing arguments
mock_parse_args.return_value = MagicMock()
def test_main_function(self, mock_send_chat, mock_evaluate, mock_parse_args, mock_asg_scheduler, mock_ec2_scheduler, mock_aurora_scheduler, mock_rds_scheduler, mock_ecs_scheduler, mock_boto3_session):
# Setup mock for boto3.Session
mock_session = MagicMock()
mock_boto3_session.return_value = mock_session

# Mock parse_arguments to return the required arguments
mock_parse_args.return_value = argparse.Namespace(
file='example/config.yml',
workspace='stage',
resource='asg,ec2,ecs,rds,aurora',
action='stop',
no_wait=False,
threads=10
)

# Mock evaluation of config
mock_evaluate.return_value = (
{"aws_region": "us-west-2", "asg": {"name": ["asg1"]}, "ec2": {"name": ["ec2-1"]},
"ecs": {"name": ["ecs1"]}, "rds": {"name": ["rds1"]}},
["ec2", "asg", "ecs", "rds"],
"stop",
"workspace_name",
False,
10
{
"aws_region": "us-west-2",
"asg": {"name": ["asg1"]},
"ec2": {"name": ["ec2-1"]},
"ecs": {"name": ["ecs1"]},
"rds": {"name": ["rds1"]},
"aurora": {"name": ["aurora1"]},
"notification": {
"enable": True,
"platform": "google",
"webhook_url": "http://example.com/webhook"
}
},
["asg", "ec2", "ecs", "rds", "aurora"],
"stop"
)

# Mock AWS session login
mock_session = MagicMock()
mock_aws_login.return_value = mock_session

# Mock scheduler functions
mock_asg_scheduler.return_value = ["asg1 stopped"]
mock_asg_scheduler.return_value = [{"ASGName": "asg1", "Status": "stopped"}]
mock_ec2_scheduler.return_value = [{"InstanceId": "ec2-1", "State": "stopped"}]
mock_ecs_scheduler.return_value = ["ecs1 stopped"]
mock_rds_scheduler.return_value = ["rds1 stopped"]
mock_ecs_scheduler.return_value = [{"ServiceName": "ecs1", "Status": "stopped"}]
mock_rds_scheduler.return_value = [{"DBInstanceIdentifier": "rds1", "DBInstanceStatus": "stopped"}]
mock_aurora_scheduler.return_value = [{"DBClusterIdentifier": "aurora1", "Status": "stopped"}]

# Force an exception to simulate an error during processing
mock_asg_scheduler.side_effect = Exception("Simulated failure")

# Run the main function
main()

# Assertions for EC2
mock_ec2_scheduler.assert_called_once_with(
{"name": ["ec2-1"]},
"stop",
instance_attributes=['InstanceId', 'InstanceType', 'State', 'PrivateIpAddress', 'PublicIpAddress']
)

# Assertions for ASG
mock_asg_scheduler.assert_called_once_with({"name": ["asg1"]}, "stop")

# Assertions for ECS
mock_ecs_scheduler.assert_called_once_with({"name": ["ecs1"]}, "stop")

# Assertions for RDS
mock_rds_scheduler.assert_called_once_with({"name": ["rds1"]}, "stop", instance_attributes=['DBInstanceIdentifier', 'DBInstanceStatus', 'DBInstanceClass', 'Engine', 'Endpoint'])

# Verify the notification was sent
mock_send_chat.assert_called()

@patch('aws_resource_scheduler.utils.ecs.EcsModule.start_ecs_service')
@patch('aws_resource_scheduler.utils.ecs.EcsModule.stop_ecs_service')
@patch('aws_resource_scheduler.utils.ecs.EcsModule.get_ecs_service_status')
def test_safe_execution(self, mock_get_status, mock_stop, mock_start):
# Mock the behavior of get_ecs_service_status to simulate different scenarios
mock_get_status.return_value = {
"cluster_name": "test-cluster",
"service_name": "test-service",
"desired_count": 1,
"running_count": 1,
"status": "ACTIVE",
"launch_type": "FARGATE",
"task_definition": "test-task-def"
}

# Mock stop_ecs_service to raise an exception
mock_stop.side_effect = Exception("Test Exception")

# Initialize EcsModule with mocked data
session = MagicMock()
storage = MagicMock()
ecs_module = EcsModule(session, storage, "workspace", no_wait=False, threads=2)

# Call safe_execution with stop_ecs_service and service_data
service_data = {
"cluster_name": "test-cluster",
"service_name": "test-service"
}
ecs_module.safe_execution(ecs_module.stop_ecs_service, service_data)

# Verify that the error was logged and the message was appended to the scheduler_summary_message
self.assertIn("Failed to process service test-service", ecs_module.scheduler_summary_message)
# Assertions
self.assertTrue(mock_asg_scheduler.called, "ASG scheduler should have been called")
self.assertTrue(mock_ec2_scheduler.called, "EC2 scheduler should have been called")
self.assertTrue(mock_ecs_scheduler.called, "ECS scheduler should have been called")
self.assertTrue(mock_rds_scheduler.called, "RDS scheduler should have been called")
self.assertTrue(mock_aurora_scheduler.called, "Aurora scheduler should have been called")

if __name__ == '__main__':
unittest.main()

0 comments on commit 3e13a5f

Please sign in to comment.