#!/usr/bin/python3
import os
import sys
import time
import libvirt
import unittest

class test_ipxe(unittest.TestCase):

    def setUp(self):
        self.arch = os.uname()[4]
        self.conn = libvirt.open(None)
        self.network = {}
        self.domain = {}

    def tearDown(self):
        for (mode, domain) in self.domain.items():
            domain.destroy()
        for (mode, network) in self.network.items():
            network.destroy()
        self.conn.close()

    def do_wait(self, stream, substr, timeout, verbose = False):
        start = time.time()
        log = b''
        while time.time() - start < timeout:
            data = stream.recv(128)
            if type(data) is int and data == -2:
                continue
            log += data
            if substr in log:
                secs = time.time() - start
                print(f'{substr.decode()} ({secs:.1f}s) ... ', end = '', flush = True)
                return
            time.sleep(0.01)
        if verbose:
            print(log.decode())
        self.fail(f'not found: {substr}')
        
    def do_ipxe(self, mode):
        name = f'test-ipxe-{mode}-{self.arch}'

        network = self.conn.networkLookupByName(name)
        if not network.isActive():
            network.create()
            self.network[mode] = network

        domain = self.conn.lookupByName(name)
        domain.create()
        self.domain[mode] = domain

        stream = self.conn.newStream()
        domain.openConsole(None, stream, 0)

        self.do_wait(stream, b'Start PXE',   20)
        self.do_wait(stream, b'iPXE',        40)
        self.do_wait(stream, b'hello world', 90)

    def test_ipv4(self):
        self.do_ipxe('ipv4')

    def test_ipv6(self):
        self.do_ipxe('ipv6')

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