def poke(self, context):
ec2_hook = EC2Hook(aws_conn_id=self.aws_conn_id,
region_name=self.region_name)
instance_state = ec2_hook.get_instance_state(
instance_id=self.instance_id)
self.log.info("instance state: %s", instance_state)
return instance_state == self.target_state
def test_init(self):
ec2_hook = EC2Hook(
aws_conn_id="aws_conn_test",
region_name="region-test",
)
self.assertEqual(ec2_hook.aws_conn_id, "aws_conn_test")
self.assertEqual(ec2_hook.region_name, "region-test")
def test_init(self):
ec2_hook = EC2Hook(
aws_conn_id="aws_conn_test",
region_name="region-test",
)
assert ec2_hook.aws_conn_id == "aws_conn_test"
assert ec2_hook.region_name == "region-test"
def execute(self, context):
ec2_hook = EC2Hook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
self.log.info("Stopping EC2 instance %s", self.instance_id)
instance = ec2_hook.get_instance(instance_id=self.instance_id)
instance.stop()
ec2_hook.wait_for_state(
instance_id=self.instance_id, target_state="stopped", check_interval=self.check_interval,
)
def test_get_instance(self):
ec2_hook = EC2Hook()
created_instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
created_instance_id = created_instances[0].instance_id
# test get_instance method
existing_instance = ec2_hook.get_instance(
instance_id=created_instance_id)
self.assertEqual(created_instance_id, existing_instance.instance_id)
def test_get_instance_state(self):
ec2_hook = EC2Hook()
created_instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
created_instance_id = created_instances[0].instance_id
all_instances = list(ec2_hook.conn.instances.all())
created_instance_state = all_instances[0].state["Name"]
# test get_instance_state method
existing_instance_state = ec2_hook.get_instance_state(
instance_id=created_instance_id)
self.assertEqual(created_instance_state, existing_instance_state)
def test_stop_instance(self):
# create instance
ec2_hook = EC2Hook()
instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
instance_id = instances[0].instance_id
# stop instance
stop_test = EC2StopInstanceOperator(
task_id="stop_test",
instance_id=instance_id,
)
stop_test.execute(None)
# assert instance state is running
self.assertEqual(ec2_hook.get_instance_state(instance_id=instance_id), "stopped")
def test_start_instance(self):
# create instance
ec2_hook = EC2Hook()
instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
instance_id = instances[0].instance_id
# start instance
start_test = EC2StartInstanceOperator(
task_id="start_test",
instance_id=instance_id,
)
start_test.execute(None)
# assert instance state is running
assert ec2_hook.get_instance_state(
instance_id=instance_id) == "running"
def test_stopped(self):
# create instance
ec2_hook = EC2Hook()
instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
instance_id = instances[0].instance_id
# start instance
ec2_hook.get_instance(instance_id=instance_id).start()
# stop sensor, waits until ec2 instance state became stopped
stop_sensor = EC2InstanceStateSensor(
task_id="stop_sensor",
target_state="stopped",
instance_id=instance_id,
)
# assert instance state is not stopped
assert not stop_sensor.poke(None)
# stop instance
ec2_hook.get_instance(instance_id=instance_id).stop()
# assert instance state is stopped
assert stop_sensor.poke(None)
def test_running(self):
# create instance
ec2_hook = EC2Hook()
instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
instance_id = instances[0].instance_id
# stop instance
ec2_hook.get_instance(instance_id=instance_id).stop()
# start sensor, waits until ec2 instance state became running
start_sensor = EC2InstanceStateSensor(
task_id="start_sensor",
target_state="running",
instance_id=instance_id,
)
# assert instance state is not running
self.assertFalse(start_sensor.poke(None))
# start instance
ec2_hook.get_instance(instance_id=instance_id).start()
# assert instance state is running
self.assertTrue(start_sensor.poke(None))
def test_terminated(self):
# create instance
ec2_hook = EC2Hook()
instances = ec2_hook.conn.create_instances(
MaxCount=1,
MinCount=1,
)
instance_id = instances[0].instance_id
# start instance
ec2_hook.get_instance(instance_id=instance_id).start()
# stop sensor, waits until ec2 instance state became terminated
stop_sensor = EC2InstanceStateSensor(
task_id="stop_sensor",
target_state="terminated",
instance_id=instance_id,
)
# assert instance state is not terminated
self.assertFalse(stop_sensor.poke(None))
# stop instance
ec2_hook.get_instance(instance_id=instance_id).terminate()
# assert instance state is terminated
self.assertTrue(stop_sensor.poke(None))
def test_get_conn_returns_boto3_resource(self):
ec2_hook = EC2Hook()
instances = list(ec2_hook.conn.instances.all())
self.assertIsNotNone(instances)