test_ec2.py

import pulumi


class MyMocks(pulumi.runtime.Mocks):
    def new_resource(self, args: pulumi.runtime.MockResourceArgs):
        outputs = args.inputs
        if args.typ == "aws:ec2/instance:Instance":
            outputs = {
                **args.inputs,
                "publicIp": "203.0.113.12",
                "publicDns": "ec2-203-0-113-12.compute-1.amazonaws.com",
            }
        return [args.name + '_id', outputs]

    def call(self, args: pulumi.runtime.MockCallArgs):
        if args.token == "aws:ec2/getAmi:getAmi":
            return {
                "architecture": "x86_64",
                "id": "ami-0eb1f3cdeeb8eed2a",
            }
        return {}


pulumi.runtime.set_mocks(MyMocks())

# Now actually import the code that creates resources, and then test it.
import infra


# Test if the service has tags and a name tag.
@pulumi.runtime.test
def test_server_tags():
    def check_tags(args):
        urn, tags = args
        assert tags, f'server {urn} must have tags'
        assert 'Name' in tags, 'server {urn} must have a name tag'

    return pulumi.Output.all(infra.server.urn, infra.server.tags).apply(check_tags)


# Test if the instance is configured with user_data.
@pulumi.runtime.test
def test_server_userdata():
    def check_user_data(args):
        urn, user_data = args
        assert user_data is None, f'illegal use of user_data on server {urn}'

    return pulumi.Output.all(infra.server.urn, infra.server.user_data).apply(check_user_data)


# Test if port 22 for ssh is exposed.
@pulumi.runtime.test
def test_security_group_rules():
    def check_security_group_rules(args):
        urn, ingress = args
        ssh_open = any(
            [rule['from_port'] == 22 and any([block == "0.0.0.0/0" for block in rule['cidr_blocks']]) for rule in
             ingress])
        assert ssh_open is False, f'security group {urn} exposes port 22 to the Internet (CIDR 0.0.0.0/0)'

    # Return the results of the unit tests.
    return pulumi.Output.all(infra.group.urn, infra.group.ingress).apply(check_security_group_rules)