Best Python code snippet using lisa_python
features.py
Source:features.py  
...17    from .platform_ import AwsPlatform18from .. import AWS19from .common import AwsNodeSchema, get_node_context20class AwsFeatureMixin:21    def _initialize_information(self, node: Node) -> None:22        node_context = get_node_context(node)23        self._instance_id = node_context.instance_id24class StartStop(AwsFeatureMixin, features.StartStop):25    def _stop(26        self,27        wait: bool = True,28        state: features.StopState = features.StopState.Shutdown,29    ) -> None:30        ec2_resource = boto3.resource("ec2")31        instance = ec2_resource.Instance(self._instance_id)32        if state == features.StopState.Hibernate:33            instance.stop(Hibernate=True)34        else:35            instance.stop()36        if wait:37            instance.wait_until_stopped()38    def _start(self, wait: bool = True) -> None:39        ec2_resource = boto3.resource("ec2")40        instance = ec2_resource.Instance(self._instance_id)41        instance.start()42        if wait:43            instance.wait_until_running()44    def _restart(self, wait: bool = True) -> None:45        ec2_resource = boto3.resource("ec2")46        instance = ec2_resource.Instance(self._instance_id)47        instance.reboot()48        if wait:49            instance.wait_until_running()50    def _initialize(self, *args: Any, **kwargs: Any) -> None:51        super()._initialize(*args, **kwargs)52        self._initialize_information(self._node)53def get_aws_disk_type(disk_type: schema.DiskType) -> str:54    assert isinstance(disk_type, schema.DiskType), (55        f"the disk_type must be one value when calling get_disk_type. "56        f"But it's {disk_type}"57    )58    result = _disk_type_mapping.get(disk_type, None)59    assert result, f"unknown disk type: {disk_type}"60    return result61# There are more disk types in AWS than Azure, like io2/gp3/io 2 Block Express.62# If need to verify the storage performance of other types, please update the mapping.63# Refer to https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ebs-volume-types.html.64# DiskType.Ephemeral is not supported on AWS now.65_disk_type_mapping: Dict[schema.DiskType, str] = {66    schema.DiskType.PremiumSSDLRS: "io1",67    schema.DiskType.StandardHDDLRS: "st1",68    schema.DiskType.StandardSSDLRS: "gp3",69}70# Tuple: (IOPS, Disk Size)71_disk_size_iops_map: Dict[schema.DiskType, List[Tuple[int, int]]] = {72    schema.DiskType.PremiumSSDLRS: [73        (100, 4),74        (1000, 64),75        (5000, 128),76        (10000, 256),77        (20000, 1024),78        (32000, 2048),79        (50000, 8192),80        (64000, 16384),81    ],82    schema.DiskType.StandardHDDLRS: [83        (100, 125),84        (200, 8192),85        (500, 16384),86    ],87    schema.DiskType.StandardSSDLRS: [88        (3000, 1),89        (6000, 128),90        (8000, 256),91        (10000, 512),92        (12000, 1024),93        (14000, 4096),94        (16000, 16384),95    ],96}97class SerialConsole(AwsFeatureMixin, features.SerialConsole):98    def _initialize(self, *args: Any, **kwargs: Any) -> None:99        super()._initialize(*args, **kwargs)100        self._initialize_information(self._node)101    def _get_console_log(self, saved_path: Optional[Path]) -> bytes:102        platform: AwsPlatform = self._platform  # type: ignore103        ec2_client = platform._ec2_client104        if saved_path:105            screenshot_response = ec2_client.get_console_screenshot(106                InstanceId=self._instance_id107            )108            screenshot_name = saved_path.joinpath("serial_console.jpg")109            with open(screenshot_name, "wb") as f:110                f.write(111                    base64.decodebytes(screenshot_response["ImageData"].encode("utf-8"))112                )113        diagnostic_data = ec2_client.get_console_output(InstanceId=self._instance_id)114        output_bytes = diagnostic_data["Output"].encode("ascii")115        return base64.b64decode(output_bytes)116class NetworkInterface(AwsFeatureMixin, features.NetworkInterface):117    """118    This Network interface feature is mainly to associate Aws119    network interface options settings.120    """121    @classmethod122    def settings_type(cls) -> Type[schema.FeatureSettings]:123        return schema.NetworkInterfaceOptionSettings124    def _initialize(self, *args: Any, **kwargs: Any) -> None:125        super()._initialize(*args, **kwargs)126        self._initialize_information(self._node)127    def _get_primary(self, nics: List[Any]) -> Any:128        found_primary = False129        for nic in nics:130            if nic["Attachment"]["DeviceIndex"] == 0:131                found_primary = True132                break133        if not found_primary:134            raise LisaException(f"fail to find primary nic for vm {self._node.name}")135        return nic136    def switch_sriov(137        self, enable: bool, wait: bool = True, reset_connections: bool = True138    ) -> None:139        aws_platform: AwsPlatform = self._platform  # type: ignore140        instance = boto3.resource("ec2").Instance(self._instance_id)141        # Don't check Intel 82599 Virtual Function (VF) interface at current142        if instance.ena_support == enable:143            self._log.debug(144                f"The accelerated networking default "145                f"status [{instance.ena_support}] is "146                f"consistent with set status [{enable}], no need to update."147            )148        else:149            self._log.debug(150                f"The accelerated networking default "151                f"status [{instance.ena_support}], "152                f"now set its status into [{enable}]."153            )154            aws_platform._ec2_client.modify_instance_attribute(155                InstanceId=instance.id,156                EnaSupport={157                    "Value": enable,158                },159            )160            instance.reload()161            assert_that(instance.ena_support).described_as(162                f"fail to set network interface accelerated "163                f"networking into status [{enable}]"164            ).is_equal_to(enable)165    def is_enabled_sriov(self) -> bool:166        instance = boto3.resource("ec2").Instance(self._instance_id)167        return instance.ena_support168    def attach_nics(169        self, extra_nic_count: int, enable_accelerated_networking: bool = True170    ) -> None:171        aws_platform: AwsPlatform = self._platform  # type: ignore172        ec2_resource = boto3.resource("ec2")173        instance = ec2_resource.Instance(self._instance_id)174        current_nic_count = len(instance.network_interfaces)175        nic_count_after_add_extra = extra_nic_count + current_nic_count176        assert (177            self._node.capability.network_interface178            and self._node.capability.network_interface.max_nic_count179        )180        assert isinstance(181            self._node.capability.network_interface.max_nic_count, int182        ), f"actual: {type(self._node.capability.network_interface.max_nic_count)}"183        node_capability_nic_count = (184            self._node.capability.network_interface.max_nic_count185        )186        if nic_count_after_add_extra > node_capability_nic_count:187            raise LisaException(188                f"nic count after add extra nics is {nic_count_after_add_extra},"189                f" it exceeds the vm size's capability {node_capability_nic_count}."190            )191        nic = self._get_primary(instance.network_interfaces_attribute)192        index = current_nic_count193        while index < current_nic_count + extra_nic_count - 1:194            extra_nic_name = f"{self._node.name}-extra-{index}"195            self._log.debug(f"start to create the nic {extra_nic_name}.")196            network_interface = ec2_resource.create_network_interface(197                Description=extra_nic_name,198                Groups=[199                    nic["Groups"][0]["GroupId"],200                ],201                SubnetId=nic["SubnetId"],202            )203            self._log.debug(204                f"start to attach the nic {extra_nic_name} into VM {self._node.name}."205            )206            aws_platform._ec2_client.attach_network_interface(207                DeviceIndex=index,208                InstanceId=instance.id,209                NetworkInterfaceId=network_interface.network_interface_id,210            )211            self._log.debug(212                f"attach the nic {extra_nic_name} into"213                f"VM {self._node.name} successfully."214            )215            index += 1216    def remove_extra_nics(self) -> None:217        aws_platform: AwsPlatform = self._platform  # type: ignore218        instance = boto3.resource("ec2").Instance(self._instance_id)219        for network_interface in instance.network_interfaces_attribute:220            if network_interface["Attachment"]["DeviceIndex"] != 0:221                aws_platform._ec2_client.detach_network_interface(222                    AttachmentId=network_interface["Attachment"]["AttachmentId"]223                )224                aws_platform._ec2_client.delete_network_interface(225                    NetworkInterfaceId=network_interface["NetworkInterfaceId"]226                )227        nic = self._get_primary(instance.network_interfaces_attribute)228        networkinterface_id: str = nic["NetworkInterfaceId"]229        self._log.debug(230            f"Only associated nic {networkinterface_id} into VM {self._node.name}."231        )232# TODO: GPU feature is not verified yet.233class Gpu(AwsFeatureMixin, features.Gpu):234    # Only contains some types of ec2 instances which support GPU here.235    # Please refer to the following link for more types:236    # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-nvidia-driver.html237    grid_supported_skus = [238        "g3s.xlarge",239        "g3.4xlarge",240        "g4dn.xlarge",241        "g5.xlarge",242        "g5.2xlarge",243    ]244    cuda_supported_skus = [245        "g3s.xlarge",246        "g3.4xlarge",247        "g4dn.xlarge",248        "g5.xlarge",249        "g5.2xlarge",250        "p2.xlarge",251        "p3.2xlarge",252    ]253    def is_supported(self) -> bool:254        # TODO: more supportability can be defined here255        supported_distro = (CentOs, Redhat, Ubuntu, Suse)256        if isinstance(self._node.os, supported_distro):257            return True258        return False259    def _initialize(self, *args: Any, **kwargs: Any) -> None:260        super()._initialize(*args, **kwargs)261        self._initialize_information(self._node)262    def _get_supported_driver(self) -> List[ComputeSDK]:263        driver_list = []264        node_runbook = self._node.capability.get_extended_runbook(AwsNodeSchema, AWS)265        if any(map((node_runbook.vm_size).__contains__, self.grid_supported_skus)):266            driver_list.append(ComputeSDK.GRID)267        if any(map((node_runbook.vm_size).__contains__, self.cuda_supported_skus)):268            driver_list.append(ComputeSDK.CUDA)269        if not driver_list:270            raise LisaException(271                "No valid Compute SDK found to install for the VM size -"272                f" {node_runbook.vm_size}."273            )274        return driver_list275class Disk(AwsFeatureMixin, features.Disk):276    """277    This Disk feature is mainly to associate Aws disk options settings.278    """279    @classmethod280    def settings_type(cls) -> Type[schema.FeatureSettings]:281        return AwsDiskOptionSettings282    def _initialize(self, *args: Any, **kwargs: Any) -> None:283        super()._initialize(*args, **kwargs)284        self._initialize_information(self._node)285    def get_raw_data_disks(self) -> List[str]:286        # Return all EBS devices except the root device287        instance = boto3.resource("ec2").Instance(self._instance_id)288        disk_array: List[str] = []289        for device_mapping in instance.block_device_mappings:290            if "Ebs" in device_mapping and "DeviceName" in device_mapping:291                if device_mapping["DeviceName"] != instance.root_device_name:292                    disk_array.append(device_mapping["DeviceName"])293        return disk_array294@dataclass_json()295@dataclass()296class AwsDiskOptionSettings(schema.DiskOptionSettings):297    def __eq__(self, o: object) -> bool:298        assert isinstance(o, AwsDiskOptionSettings), f"actual: {type(o)}"...Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!
