#!/usr/bin/python3 -cimport os, sys; os.execv(os.path.dirname(sys.argv[1]) + "/../common/pywrap", sys.argv)## Copyright (C) 2018 Red Hat, Inc.# SPDX-License-Identifier: LGPL-2.1-or-laterimport jsonimport osimport shutilimport subprocessimport tempfileimport unittestfrom glob import globimport testlibfrom machine import testvm@unittest.skipUnless("TEST_OS" in os.environ, "TEST_OS not set")class TestImageCustomize(unittest.TestCase):def checkBoot(self, image):with testvm.Timeout(seconds=300, error_message="Timed out waiting for image to run"):network = testvm.VirtNetwork(0, image=image)machine = testvm.VirtMachine(image=image, networking=network.host())machine.boot()out = machine.execute('cat /var/custom-test')machine.stop()self.assertEqual(out, "hello\n")def testCustomDir(self):dest = tempfile.mkdtemp()self.addCleanup(shutil.rmtree, dest)img = os.path.join(dest, os.environ["TEST_OS"])with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose", "--run-command","echo hello > /var/custom-test", img])self.assertTrue(os.path.exists(img))self.checkBoot(img)def testBaseImage(self):img = "custom-" + os.environ["TEST_OS"]def cleanup():for f in glob(f"test/images/{img}*"):os.unlink(f)self.addCleanup(cleanup)with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose", "--run-command","echo hello > /var/custom-test", "--base-image", os.environ["TEST_OS"], img])self.assertTrue(os.path.exists(os.path.join("test/images", img)))# notice, not giving directory here - test/images/ should be the defaultself.checkBoot(img)def testScriptRelativePath(self):dest = tempfile.mkdtemp(dir=".")self.addCleanup(shutil.rmtree, dest)script = os.path.join(dest, "setup.sh")with open(script, "w") as f:f.write("#!/bin/sh -eu\necho hello > /var/custom-test\n")img = os.path.join(dest, os.environ["TEST_OS"])with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose", "--script", script, img])self.assertTrue(os.path.exists(img))self.checkBoot(img)def testUpload(self):dest = tempfile.mkdtemp(dir=".")self.addCleanup(shutil.rmtree, dest)img = os.path.join(dest, os.environ["TEST_OS"])with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose","--upload", "/etc/passwd:/tmp/passwd","--run-command", "echo hello > /var/custom-test","--run-command", "grep ^root: /tmp/passwd", img])self.checkBoot(img)def testFailurePropagation(self):dest = tempfile.mkdtemp(dir=".")self.addCleanup(shutil.rmtree, dest)img = os.path.join(dest, os.environ["TEST_OS"])with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose","--run-command", "true", img])with self.assertRaises(subprocess.CalledProcessError):subprocess.check_call(["bots/image-customize", "--verbose","--run-command", "false", img])def testResize(self):dest = tempfile.mkdtemp(dir=".")self.addCleanup(shutil.rmtree, dest)img = os.path.join(dest, os.environ["TEST_OS"])with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose","--resize", "30G", img])output = subprocess.check_output(["qemu-img", "info", "--output=json", img], encoding="utf-8")info = json.loads(output)self.assertEqual(int(info['virtual-size']) // 1024 // 1024 // 1024, 30)@unittest.skipUnless("TEST_OS" in os.environ, "TEST_OS not set")class TestBotsVM(unittest.TestCase):def testBasic(self):dest = tempfile.mkdtemp()self.addCleanup(shutil.rmtree, dest)img = os.path.join(dest, os.environ["TEST_OS"])with testvm.Timeout(seconds=300, error_message="Timed out waiting for image-customize"):subprocess.check_call(["bots/image-customize", "--verbose", "--run-command","echo hello > /var/custom-test", img])# boot it and wait for RUNNING marker, parse out ssh and cockpit addresseswith testvm.Timeout(seconds=300, error_message="Timed out waiting for testvm.py to boot VM"):vm = subprocess.Popen(["bots/machine/testvm.py", img],stdout=subprocess.PIPE, universal_newlines=True)# first line should be the SSH commandssh_command = vm.stdout.readline().split()# second line is the redirected cockpit addresscockpit_address = vm.stdout.readline()# third should be the "I am ready" flagrunning = vm.stdout.readline()self.assertEqual(running, "RUNNING\n")self.assertTrue(cockpit_address.startswith("http://127.0.0.2:9"), cockpit_address)# test SSH command and that we have the expected flag fileself.assertEqual(ssh_command[0], "ssh")with testvm.Timeout(seconds=30, error_message="Timed out waiting for ssh command"):out = subprocess.check_output([*ssh_command, "cat", "/var/custom-test"])self.assertEqual(out, b"hello\n")# should cleanly stop on SIGTERMvm.terminate()with testvm.Timeout(seconds=60, error_message="Timed out waiting for script to terminate"):self.assertEqual(vm.wait(), 0)if __name__ == '__main__':testlib.test_main()