# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Tests that ensure the boot time to init process is within spec.

The boot time tests use the minimal kernel image and root file system found in
the S3 bucket at https://s3.amazonaws.com/spec.ccfc.min/img/minimal.
For boot time testing the serial console is not enabled.
"""

import re
import time
import platform

from framework import decorators

# The maximum acceptable boot time in us.
MAX_BOOT_TIME_US = 150000
# NOTE: for aarch64 most of the boot time is spent by the kernel to unpack the
# initramfs in RAM. This time is influenced by the size and the compression
# method of the used initrd image.
INITRD_BOOT_TIME_US = 160000 if platform.machine() == "x86_64" else 500000
# TODO: Keep a `current` boot time in S3 and validate we don't regress
# Regex for obtaining boot time from some string.
TIMESTAMP_LOG_REGEX = r'Guest-boot-time\s+\=\s+(\d+)\s+us'

NO_OF_MICROVMS = 10


def test_single_microvm_boottime_no_network(test_microvm_with_boottime):
    """Check guest boottime of microvm without network."""
    _ = _configure_vm(test_microvm_with_boottime)
    time.sleep(0.4)
    boottime_us = _test_microvm_boottime(test_microvm_with_boottime.log_data)
    print("Boot time with no network is: " + str(boottime_us) + " us")


@decorators.test_context('boottime', NO_OF_MICROVMS)
def test_multiple_microvm_boottime_no_network(test_multiple_microvms):
    """Check guest boottime without network when spawning multiple microvms."""
    microvms = test_multiple_microvms
    assert microvms
    assert len(microvms) == NO_OF_MICROVMS
    log_fifos_data = []
    for i in range(NO_OF_MICROVMS):
        _ = _configure_vm(microvms[i])
        time.sleep(0.4)
        log_fifos_data.append(microvms[i].log_data)
    for i in range(NO_OF_MICROVMS):
        _ = _test_microvm_boottime(log_fifos_data[i])


@decorators.test_context('boottime', NO_OF_MICROVMS)
def test_multiple_microvm_boottime_with_network(
        test_multiple_microvms,
        network_config
):
    """Check guest boottime with network when spawning multiple microvms."""
    microvms = test_multiple_microvms
    assert microvms
    assert len(microvms) == NO_OF_MICROVMS
    log_fifos_data = []
    _taps = []
    for i in range(NO_OF_MICROVMS):
        _tap = _configure_vm(microvms[i], {
            "config": network_config, "iface_id": str(i)
        })
        time.sleep(0.4)
        log_fifos_data.append(microvms[i].log_data)
        _taps.append(_tap)
    for i in range(NO_OF_MICROVMS):
        _ = _test_microvm_boottime(log_fifos_data[i])


def test_single_microvm_boottime_with_network(
        test_microvm_with_boottime,
        network_config
):
    """Check guest boottime of microvm with network."""
    _tap = _configure_vm(test_microvm_with_boottime, {
        "config": network_config, "iface_id": "1"
    })
    time.sleep(0.4)
    boottime_us = _test_microvm_boottime(test_microvm_with_boottime.log_data)
    print("Boot time with network configured is: " + str(boottime_us) + " us")


def test_single_microvm_initrd_boottime(
        test_microvm_with_initrd):
    """Check guest boottime of microvm with initrd."""
    _tap = _configure_vm(test_microvm_with_initrd, initrd=True)
    time.sleep(0.8)
    boottime_us = _test_microvm_boottime(
        test_microvm_with_initrd.log_data, max_time_us=INITRD_BOOT_TIME_US)
    print("Boot time with initrd is: " + str(boottime_us) + " us")


def _test_microvm_boottime(log_fifo_data, max_time_us=MAX_BOOT_TIME_US):
    """Assert that we meet the minimum boot time.

    TODO: Should use a microVM with the `boottime` capability.
    """
    boot_time_us = 0
    timestamps = re.findall(TIMESTAMP_LOG_REGEX, log_fifo_data)
    if timestamps:
        boot_time_us = int(timestamps[0])

    assert boot_time_us > 0
    assert boot_time_us < max_time_us
    return boot_time_us


def _configure_vm(microvm, network_info=None, initrd=False):
    """Auxiliary function for preparing microvm before measuring boottime."""
    microvm.spawn()

    # Machine configuration specified in the SLA.
    config = {
        'vcpu_count': 1,
        'mem_size_mib': 128
    }

    if initrd:
        config['add_root_device'] = False
        config['use_initrd'] = True

    microvm.basic_config(**config)

    if network_info:
        _tap, _, _ = microvm.ssh_network_config(
            network_info["config"],
            network_info["iface_id"]
        )

    microvm.start()
    return _tap if network_info else None

@@ -426,42 +426,8 @@
 rm%0A%0A
-from framework import decorators%0A%0A
 # Th
@@ -940,53 +940,17 @@
 s'%0A%0A
-NO_OF_MICROVMS = 10%0A%0A%0Adef test_single_microvm
+%0Adef test
 _boo
@@ -1279,1323 +1279,16 @@
 )%0A%0A%0A
-@decorators.test_context('boottime', NO_OF_MICROVMS)%0Adef test_multiple_microvm_boottime_no_network(test_multiple_microvms):%0A    %22%22%22Check guest boottime without network when spawning multiple microvms.%22%22%22%0A    microvms = test_multiple_microvms%0A    assert microvms%0A    assert len(microvms) == NO_OF_MICROVMS%0A    log_fifos_data = %5B%5D%0A    for i in range(NO_OF_MICROVMS):%0A        _ = _configure_vm(microvms%5Bi%5D)%0A        time.sleep(0.4)%0A        log_fifos_data.append(microvms%5Bi%5D.log_data)%0A    for i in range(NO_OF_MICROVMS):%0A        _ = _test_microvm_boottime(log_fifos_data%5Bi%5D)%0A%0A%0A@decorators.test_context('boottime', NO_OF_MICROVMS)%0Adef test_multiple_microvm_boottime_with_network(%0A        test_multiple_microvms,%0A        network_config%0A):%0A    %22%22%22Check guest boottime with network when spawning multiple microvms.%22%22%22%0A    microvms = test_multiple_microvms%0A    assert microvms%0A    assert len(microvms) == NO_OF_MICROVMS%0A    log_fifos_data = %5B%5D%0A    _taps = %5B%5D%0A    for i in range(NO_OF_MICROVMS):%0A        _tap = _configure_vm(microvms%5Bi%5D, %7B%0A            %22config%22: network_config, %22iface_id%22: str(i)%0A        %7D)%0A        time.sleep(0.4)%0A        log_fifos_data.append(microvms%5Bi%5D.log_data)%0A        _taps.append(_tap)%0A    for i in range(NO_OF_MICROVMS):%0A        _ = _test_microvm_boottime(log_fifos_data%5Bi%5D)%0A%0A%0Adef test_single_microvm
+def test
 _boo
@@ -1729,23 +1729,8 @@
 est_
-single_microvm_
 init

remove irrelevant performance test