Synnefo » snf-ganeti

#!/usr/bin/python

#

# Copyright (C) 2006, 2007 Google Inc.

#

# This program is free software; you can redistribute it and/or modify

# it under the terms of the GNU General Public License as published by

# the Free Software Foundation; either version 2 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful, but

# WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

# General Public License for more details.

#

# You should have received a copy of the GNU General Public License

# along with this program; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA

# 02110-1301, USA.

"""Script for unittesting the utils module"""

import unittest

import os

import time

import tempfile

import os.path

import os

import stat

import signal

import socket

import shutil

import re

import select

import string

import OpenSSL

import warnings

import distutils.version

import glob

import ganeti

import testutils

from ganeti import constants

from ganeti import utils

from ganeti import errors

from ganeti.utils import IsProcessAlive, RunCmd, \

     RemoveFile, MatchNameComponent, FormatUnit, \

     ParseUnit, AddAuthorizedKey, RemoveAuthorizedKey, \

     ShellQuote, ShellQuoteArgs, TcpPing, ListVisibleFiles, \

     SetEtcHostsEntry, RemoveEtcHostsEntry, FirstFree, OwnIpAddress, \

     TailFile, ForceDictType, SafeEncode, IsNormAbsPath, FormatTime, \

     UnescapeAndSplit, RunParts, PathJoin, HostInfo

from ganeti.errors import LockError, UnitParseError, GenericError, \

     ProgrammerError, OpPrereqError

class TestIsProcessAlive(unittest.TestCase):

  """Testing case for IsProcessAlive"""

  def testExists(self):

    mypid = os.getpid()

    self.assert_(IsProcessAlive(mypid),

                 "can't find myself running")

  def testNotExisting(self):

    pid_non_existing = os.fork()

    if pid_non_existing == 0:

      os._exit(0)

    elif pid_non_existing < 0:

      raise SystemError("can't fork")

    os.waitpid(pid_non_existing, 0)

    self.assert_(not IsProcessAlive(pid_non_existing),

                 "nonexisting process detected")

class TestPidFileFunctions(unittest.TestCase):

  """Tests for WritePidFile, RemovePidFile and ReadPidFile"""

  def setUp(self):

    self.dir = tempfile.mkdtemp()

    self.f_dpn = lambda name: os.path.join(self.dir, "%s.pid" % name)

    utils.DaemonPidFileName = self.f_dpn

  def testPidFileFunctions(self):

    pid_file = self.f_dpn('test')

    utils.WritePidFile('test')

    self.failUnless(os.path.exists(pid_file),

                    "PID file should have been created")

    read_pid = utils.ReadPidFile(pid_file)

    self.failUnlessEqual(read_pid, os.getpid())

    self.failUnless(utils.IsProcessAlive(read_pid))

    self.failUnlessRaises(GenericError, utils.WritePidFile, 'test')

    utils.RemovePidFile('test')

    self.failIf(os.path.exists(pid_file),

                "PID file should not exist anymore")

    self.failUnlessEqual(utils.ReadPidFile(pid_file), 0,

                         "ReadPidFile should return 0 for missing pid file")

    fh = open(pid_file, "w")

    fh.write("blah\n")

    fh.close()

    self.failUnlessEqual(utils.ReadPidFile(pid_file), 0,

                         "ReadPidFile should return 0 for invalid pid file")

    utils.RemovePidFile('test')

    self.failIf(os.path.exists(pid_file),

                "PID file should not exist anymore")

  def testKill(self):

    pid_file = self.f_dpn('child')

    r_fd, w_fd = os.pipe()

    new_pid = os.fork()

    if new_pid == 0: #child

      utils.WritePidFile('child')

      os.write(w_fd, 'a')

      signal.pause()

      os._exit(0)

      return

    # else we are in the parent

    # wait until the child has written the pid file

    os.read(r_fd, 1)

    read_pid = utils.ReadPidFile(pid_file)

    self.failUnlessEqual(read_pid, new_pid)

    self.failUnless(utils.IsProcessAlive(new_pid))

    utils.KillProcess(new_pid, waitpid=True)

    self.failIf(utils.IsProcessAlive(new_pid))

    utils.RemovePidFile('child')

    self.failUnlessRaises(ProgrammerError, utils.KillProcess, 0)

  def tearDown(self):

    for name in os.listdir(self.dir):

      os.unlink(os.path.join(self.dir, name))

    os.rmdir(self.dir)

class TestRunCmd(testutils.GanetiTestCase):

  """Testing case for the RunCmd function"""

  def setUp(self):

    testutils.GanetiTestCase.setUp(self)

    self.magic = time.ctime() + " ganeti test"

    self.fname = self._CreateTempFile()

  def testOk(self):

    """Test successful exit code"""

    result = RunCmd("/bin/sh -c 'exit 0'")

    self.assertEqual(result.exit_code, 0)

    self.assertEqual(result.output, "")

  def testFail(self):

    """Test fail exit code"""

    result = RunCmd("/bin/sh -c 'exit 1'")

    self.assertEqual(result.exit_code, 1)

    self.assertEqual(result.output, "")

  def testStdout(self):

    """Test standard output"""

    cmd = 'echo -n "%s"' % self.magic

    result = RunCmd("/bin/sh -c '%s'" % cmd)

    self.assertEqual(result.stdout, self.magic)

    result = RunCmd("/bin/sh -c '%s'" % cmd, output=self.fname)

    self.assertEqual(result.output, "")

    self.assertFileContent(self.fname, self.magic)

  def testStderr(self):

    """Test standard error"""

    cmd = 'echo -n "%s"' % self.magic

    result = RunCmd("/bin/sh -c '%s' 1>&2" % cmd)

    self.assertEqual(result.stderr, self.magic)

    result = RunCmd("/bin/sh -c '%s' 1>&2" % cmd, output=self.fname)

    self.assertEqual(result.output, "")

    self.assertFileContent(self.fname, self.magic)

  def testCombined(self):

    """Test combined output"""

    cmd = 'echo -n "A%s"; echo -n "B%s" 1>&2' % (self.magic, self.magic)

    expected = "A" + self.magic + "B" + self.magic

    result = RunCmd("/bin/sh -c '%s'" % cmd)

    self.assertEqual(result.output, expected)

    result = RunCmd("/bin/sh -c '%s'" % cmd, output=self.fname)

    self.assertEqual(result.output, "")

    self.assertFileContent(self.fname, expected)

  def testSignal(self):

    """Test signal"""

    result = RunCmd(["python", "-c", "import os; os.kill(os.getpid(), 15)"])

    self.assertEqual(result.signal, 15)

    self.assertEqual(result.output, "")

  def testListRun(self):

    """Test list runs"""

    result = RunCmd(["true"])

    self.assertEqual(result.signal, None)

    self.assertEqual(result.exit_code, 0)

    result = RunCmd(["/bin/sh", "-c", "exit 1"])

    self.assertEqual(result.signal, None)

    self.assertEqual(result.exit_code, 1)

    result = RunCmd(["echo", "-n", self.magic])

    self.assertEqual(result.signal, None)

    self.assertEqual(result.exit_code, 0)

    self.assertEqual(result.stdout, self.magic)

  def testFileEmptyOutput(self):

    """Test file output"""

    result = RunCmd(["true"], output=self.fname)

    self.assertEqual(result.signal, None)

    self.assertEqual(result.exit_code, 0)

    self.assertFileContent(self.fname, "")

  def testLang(self):

    """Test locale environment"""

    old_env = os.environ.copy()

    try:

      os.environ["LANG"] = "en_US.UTF-8"

      os.environ["LC_ALL"] = "en_US.UTF-8"

      result = RunCmd(["locale"])

      for line in result.output.splitlines():

        key, value = line.split("=", 1)

        # Ignore these variables, they're overridden by LC_ALL

        if key == "LANG" or key == "LANGUAGE":

          continue

        self.failIf(value and value != "C" and value != '"C"',

            "Variable %s is set to the invalid value '%s'" % (key, value))

    finally:

      os.environ = old_env

  def testDefaultCwd(self):

    """Test default working directory"""

    self.failUnlessEqual(RunCmd(["pwd"]).stdout.strip(), "/")

  def testCwd(self):

    """Test default working directory"""

    self.failUnlessEqual(RunCmd(["pwd"], cwd="/").stdout.strip(), "/")

    self.failUnlessEqual(RunCmd(["pwd"], cwd="/tmp").stdout.strip(), "/tmp")

    cwd = os.getcwd()

    self.failUnlessEqual(RunCmd(["pwd"], cwd=cwd).stdout.strip(), cwd)

  def testResetEnv(self):

    """Test environment reset functionality"""

    self.failUnlessEqual(RunCmd(["env"], reset_env=True).stdout.strip(), "")

    self.failUnlessEqual(RunCmd(["env"], reset_env=True,

                                env={"FOO": "bar",}).stdout.strip(), "FOO=bar")

class TestRunParts(unittest.TestCase):

  """Testing case for the RunParts function"""

  def setUp(self):

    self.rundir = tempfile.mkdtemp(prefix="ganeti-test", suffix=".tmp")

  def tearDown(self):

    shutil.rmtree(self.rundir)

  def testEmpty(self):

    """Test on an empty dir"""

    self.failUnlessEqual(RunParts(self.rundir, reset_env=True), [])

  def testSkipWrongName(self):

    """Test that wrong files are skipped"""

    fname = os.path.join(self.rundir, "00test.dot")

    utils.WriteFile(fname, data="")

    os.chmod(fname, stat.S_IREAD | stat.S_IEXEC)

    relname = os.path.basename(fname)

    self.failUnlessEqual(RunParts(self.rundir, reset_env=True),

                         [(relname, constants.RUNPARTS_SKIP, None)])

  def testSkipNonExec(self):

    """Test that non executable files are skipped"""

    fname = os.path.join(self.rundir, "00test")

    utils.WriteFile(fname, data="")

    relname = os.path.basename(fname)

    self.failUnlessEqual(RunParts(self.rundir, reset_env=True),

                         [(relname, constants.RUNPARTS_SKIP, None)])

  def testError(self):

    """Test error on a broken executable"""

    fname = os.path.join(self.rundir, "00test")

    utils.WriteFile(fname, data="")

    os.chmod(fname, stat.S_IREAD | stat.S_IEXEC)

    (relname, status, error) = RunParts(self.rundir, reset_env=True)[0]

    self.failUnlessEqual(relname, os.path.basename(fname))

    self.failUnlessEqual(status, constants.RUNPARTS_ERR)

    self.failUnless(error)

  def testSorted(self):

    """Test executions are sorted"""

    files = []

    files.append(os.path.join(self.rundir, "64test"))

    files.append(os.path.join(self.rundir, "00test"))

    files.append(os.path.join(self.rundir, "42test"))

    for fname in files:

      utils.WriteFile(fname, data="")

    results = RunParts(self.rundir, reset_env=True)

    for fname in sorted(files):

      self.failUnlessEqual(os.path.basename(fname), results.pop(0)[0])

  def testOk(self):

    """Test correct execution"""

    fname = os.path.join(self.rundir, "00test")

    utils.WriteFile(fname, data="#!/bin/sh\n\necho -n ciao")

    os.chmod(fname, stat.S_IREAD | stat.S_IEXEC)

    (relname, status, runresult) = RunParts(self.rundir, reset_env=True)[0]

    self.failUnlessEqual(relname, os.path.basename(fname))

    self.failUnlessEqual(status, constants.RUNPARTS_RUN)

    self.failUnlessEqual(runresult.stdout, "ciao")

  def testRunFail(self):

    """Test correct execution, with run failure"""

    fname = os.path.join(self.rundir, "00test")

    utils.WriteFile(fname, data="#!/bin/sh\n\nexit 1")

    os.chmod(fname, stat.S_IREAD | stat.S_IEXEC)

    (relname, status, runresult) = RunParts(self.rundir, reset_env=True)[0]

    self.failUnlessEqual(relname, os.path.basename(fname))

    self.failUnlessEqual(status, constants.RUNPARTS_RUN)

    self.failUnlessEqual(runresult.exit_code, 1)

    self.failUnless(runresult.failed)

  def testRunMix(self):

    files = []

    files.append(os.path.join(self.rundir, "00test"))

    files.append(os.path.join(self.rundir, "42test"))

    files.append(os.path.join(self.rundir, "64test"))

    files.append(os.path.join(self.rundir, "99test"))

    files.sort()

    # 1st has errors in execution

    utils.WriteFile(files[0], data="#!/bin/sh\n\nexit 1")

    os.chmod(files[0], stat.S_IREAD | stat.S_IEXEC)

    # 2nd is skipped

    utils.WriteFile(files[1], data="")

    # 3rd cannot execute properly

    utils.WriteFile(files[2], data="")

    os.chmod(files[2], stat.S_IREAD | stat.S_IEXEC)

    # 4th execs

    utils.WriteFile(files[3], data="#!/bin/sh\n\necho -n ciao")

    os.chmod(files[3], stat.S_IREAD | stat.S_IEXEC)

    results = RunParts(self.rundir, reset_env=True)

    (relname, status, runresult) = results[0]

    self.failUnlessEqual(relname, os.path.basename(files[0]))

    self.failUnlessEqual(status, constants.RUNPARTS_RUN)

    self.failUnlessEqual(runresult.exit_code, 1)

    self.failUnless(runresult.failed)

    (relname, status, runresult) = results[1]

    self.failUnlessEqual(relname, os.path.basename(files[1]))

    self.failUnlessEqual(status, constants.RUNPARTS_SKIP)

    self.failUnlessEqual(runresult, None)

    (relname, status, runresult) = results[2]

    self.failUnlessEqual(relname, os.path.basename(files[2]))

    self.failUnlessEqual(status, constants.RUNPARTS_ERR)

    self.failUnless(runresult)

    (relname, status, runresult) = results[3]

    self.failUnlessEqual(relname, os.path.basename(files[3]))

    self.failUnlessEqual(status, constants.RUNPARTS_RUN)

    self.failUnlessEqual(runresult.output, "ciao")

    self.failUnlessEqual(runresult.exit_code, 0)

    self.failUnless(not runresult.failed)

class TestRemoveFile(unittest.TestCase):

  """Test case for the RemoveFile function"""

  def setUp(self):

    """Create a temp dir and file for each case"""

    self.tmpdir = tempfile.mkdtemp('', 'ganeti-unittest-')

    fd, self.tmpfile = tempfile.mkstemp('', '', self.tmpdir)

    os.close(fd)

  def tearDown(self):

    if os.path.exists(self.tmpfile):

      os.unlink(self.tmpfile)

    os.rmdir(self.tmpdir)

  def testIgnoreDirs(self):

    """Test that RemoveFile() ignores directories"""

    self.assertEqual(None, RemoveFile(self.tmpdir))

  def testIgnoreNotExisting(self):

    """Test that RemoveFile() ignores non-existing files"""

    RemoveFile(self.tmpfile)

    RemoveFile(self.tmpfile)

  def testRemoveFile(self):

    """Test that RemoveFile does remove a file"""

    RemoveFile(self.tmpfile)

    if os.path.exists(self.tmpfile):

      self.fail("File '%s' not removed" % self.tmpfile)

  def testRemoveSymlink(self):

    """Test that RemoveFile does remove symlinks"""

    symlink = self.tmpdir + "/symlink"

    os.symlink("no-such-file", symlink)

    RemoveFile(symlink)

    if os.path.exists(symlink):

      self.fail("File '%s' not removed" % symlink)

    os.symlink(self.tmpfile, symlink)

    RemoveFile(symlink)

    if os.path.exists(symlink):

      self.fail("File '%s' not removed" % symlink)

class TestRename(unittest.TestCase):

  """Test case for RenameFile"""

  def setUp(self):

    """Create a temporary directory"""

    self.tmpdir = tempfile.mkdtemp()

    self.tmpfile = os.path.join(self.tmpdir, "test1")

    # Touch the file

    open(self.tmpfile, "w").close()

  def tearDown(self):

    """Remove temporary directory"""

    shutil.rmtree(self.tmpdir)

  def testSimpleRename1(self):

    """Simple rename 1"""

    utils.RenameFile(self.tmpfile, os.path.join(self.tmpdir, "xyz"))

    self.assert_(os.path.isfile(os.path.join(self.tmpdir, "xyz")))

  def testSimpleRename2(self):

    """Simple rename 2"""

    utils.RenameFile(self.tmpfile, os.path.join(self.tmpdir, "xyz"),

                     mkdir=True)

    self.assert_(os.path.isfile(os.path.join(self.tmpdir, "xyz")))

  def testRenameMkdir(self):

    """Rename with mkdir"""

    utils.RenameFile(self.tmpfile, os.path.join(self.tmpdir, "test/xyz"),

                     mkdir=True)

    self.assert_(os.path.isdir(os.path.join(self.tmpdir, "test")))

    self.assert_(os.path.isfile(os.path.join(self.tmpdir, "test/xyz")))

    utils.RenameFile(os.path.join(self.tmpdir, "test/xyz"),

                     os.path.join(self.tmpdir, "test/foo/bar/baz"),

                     mkdir=True)

    self.assert_(os.path.isdir(os.path.join(self.tmpdir, "test")))

    self.assert_(os.path.isdir(os.path.join(self.tmpdir, "test/foo/bar")))

    self.assert_(os.path.isfile(os.path.join(self.tmpdir, "test/foo/bar/baz")))

class TestMatchNameComponent(unittest.TestCase):

  """Test case for the MatchNameComponent function"""

  def testEmptyList(self):

    """Test that there is no match against an empty list"""

    self.failUnlessEqual(MatchNameComponent("", []), None)

    self.failUnlessEqual(MatchNameComponent("test", []), None)

  def testSingleMatch(self):

    """Test that a single match is performed correctly"""

    mlist = ["test1.example.com", "test2.example.com", "test3.example.com"]

    for key in "test2", "test2.example", "test2.example.com":

      self.failUnlessEqual(MatchNameComponent(key, mlist), mlist[1])

  def testMultipleMatches(self):

    """Test that a multiple match is returned as None"""

    mlist = ["test1.example.com", "test1.example.org", "test1.example.net"]

    for key in "test1", "test1.example":

      self.failUnlessEqual(MatchNameComponent(key, mlist), None)

  def testFullMatch(self):

    """Test that a full match is returned correctly"""

    key1 = "test1"

    key2 = "test1.example"

    mlist = [key2, key2 + ".com"]

    self.failUnlessEqual(MatchNameComponent(key1, mlist), None)

    self.failUnlessEqual(MatchNameComponent(key2, mlist), key2)

  def testCaseInsensitivePartialMatch(self):

    """Test for the case_insensitive keyword"""

    mlist = ["test1.example.com", "test2.example.net"]

    self.assertEqual(MatchNameComponent("test2", mlist, case_sensitive=False),

                     "test2.example.net")

    self.assertEqual(MatchNameComponent("Test2", mlist, case_sensitive=False),

                     "test2.example.net")

    self.assertEqual(MatchNameComponent("teSt2", mlist, case_sensitive=False),

                     "test2.example.net")

    self.assertEqual(MatchNameComponent("TeSt2", mlist, case_sensitive=False),

                     "test2.example.net")

  def testCaseInsensitiveFullMatch(self):

    mlist = ["ts1.ex", "ts1.ex.org", "ts2.ex", "Ts2.ex"]

    # Between the two ts1 a full string match non-case insensitive should work

    self.assertEqual(MatchNameComponent("Ts1", mlist, case_sensitive=False),

                     None)

    self.assertEqual(MatchNameComponent("Ts1.ex", mlist, case_sensitive=False),

                     "ts1.ex")

    self.assertEqual(MatchNameComponent("ts1.ex", mlist, case_sensitive=False),

                     "ts1.ex")

    # Between the two ts2 only case differs, so only case-match works

    self.assertEqual(MatchNameComponent("ts2.ex", mlist, case_sensitive=False),

                     "ts2.ex")

    self.assertEqual(MatchNameComponent("Ts2.ex", mlist, case_sensitive=False),

                     "Ts2.ex")

    self.assertEqual(MatchNameComponent("TS2.ex", mlist, case_sensitive=False),

                     None)

class TestTimestampForFilename(unittest.TestCase):

  def test(self):

    self.assert_("." not in utils.TimestampForFilename())

    self.assert_(":" not in utils.TimestampForFilename())

class TestCreateBackup(testutils.GanetiTestCase):

  def setUp(self):

    testutils.GanetiTestCase.setUp(self)

    self.tmpdir = tempfile.mkdtemp()

  def tearDown(self):

    testutils.GanetiTestCase.tearDown(self)

    shutil.rmtree(self.tmpdir)

  def testEmpty(self):

    filename = utils.PathJoin(self.tmpdir, "config.data")

    utils.WriteFile(filename, data="")

    bname = utils.CreateBackup(filename)

    self.assertFileContent(bname, "")

    self.assertEqual(len(glob.glob("%s*" % filename)), 2)

    utils.CreateBackup(filename)

    self.assertEqual(len(glob.glob("%s*" % filename)), 3)

    utils.CreateBackup(filename)

    self.assertEqual(len(glob.glob("%s*" % filename)), 4)

    fifoname = utils.PathJoin(self.tmpdir, "fifo")

    os.mkfifo(fifoname)

    self.assertRaises(errors.ProgrammerError, utils.CreateBackup, fifoname)

  def testContent(self):

    bkpcount = 0

    for data in ["", "X", "Hello World!\n" * 100, "Binary data\0\x01\x02\n"]:

      for rep in [1, 2, 10, 127]:

        testdata = data * rep

        filename = utils.PathJoin(self.tmpdir, "test.data_")

        utils.WriteFile(filename, data=testdata)

        self.assertFileContent(filename, testdata)

        for _ in range(3):

          bname = utils.CreateBackup(filename)

          bkpcount += 1

          self.assertFileContent(bname, testdata)

          self.assertEqual(len(glob.glob("%s*" % filename)), 1 + bkpcount)

class TestFormatUnit(unittest.TestCase):

  """Test case for the FormatUnit function"""

  def testMiB(self):

    self.assertEqual(FormatUnit(1, 'h'), '1M')

    self.assertEqual(FormatUnit(100, 'h'), '100M')

    self.assertEqual(FormatUnit(1023, 'h'), '1023M')

    self.assertEqual(FormatUnit(1, 'm'), '1')

    self.assertEqual(FormatUnit(100, 'm'), '100')

    self.assertEqual(FormatUnit(1023, 'm'), '1023')

    self.assertEqual(FormatUnit(1024, 'm'), '1024')

    self.assertEqual(FormatUnit(1536, 'm'), '1536')

    self.assertEqual(FormatUnit(17133, 'm'), '17133')

    self.assertEqual(FormatUnit(1024 * 1024 - 1, 'm'), '1048575')

  def testGiB(self):

    self.assertEqual(FormatUnit(1024, 'h'), '1.0G')

    self.assertEqual(FormatUnit(1536, 'h'), '1.5G')

    self.assertEqual(FormatUnit(17133, 'h'), '16.7G')

    self.assertEqual(FormatUnit(1024 * 1024 - 1, 'h'), '1024.0G')

    self.assertEqual(FormatUnit(1024, 'g'), '1.0')

    self.assertEqual(FormatUnit(1536, 'g'), '1.5')

    self.assertEqual(FormatUnit(17133, 'g'), '16.7')

    self.assertEqual(FormatUnit(1024 * 1024 - 1, 'g'), '1024.0')

    self.assertEqual(FormatUnit(1024 * 1024, 'g'), '1024.0')

    self.assertEqual(FormatUnit(5120 * 1024, 'g'), '5120.0')

    self.assertEqual(FormatUnit(29829 * 1024, 'g'), '29829.0')

  def testTiB(self):

    self.assertEqual(FormatUnit(1024 * 1024, 'h'), '1.0T')

    self.assertEqual(FormatUnit(5120 * 1024, 'h'), '5.0T')

    self.assertEqual(FormatUnit(29829 * 1024, 'h'), '29.1T')

    self.assertEqual(FormatUnit(1024 * 1024, 't'), '1.0')

    self.assertEqual(FormatUnit(5120 * 1024, 't'), '5.0')

    self.assertEqual(FormatUnit(29829 * 1024, 't'), '29.1')

class TestParseUnit(unittest.TestCase):

  """Test case for the ParseUnit function"""

  SCALES = (('', 1),

            ('M', 1), ('G', 1024), ('T', 1024 * 1024),

            ('MB', 1), ('GB', 1024), ('TB', 1024 * 1024),

            ('MiB', 1), ('GiB', 1024), ('TiB', 1024 * 1024))

  def testRounding(self):

    self.assertEqual(ParseUnit('0'), 0)

    self.assertEqual(ParseUnit('1'), 4)

    self.assertEqual(ParseUnit('2'), 4)

    self.assertEqual(ParseUnit('3'), 4)

    self.assertEqual(ParseUnit('124'), 124)

    self.assertEqual(ParseUnit('125'), 128)

    self.assertEqual(ParseUnit('126'), 128)

    self.assertEqual(ParseUnit('127'), 128)

    self.assertEqual(ParseUnit('128'), 128)

    self.assertEqual(ParseUnit('129'), 132)

    self.assertEqual(ParseUnit('130'), 132)

  def testFloating(self):

    self.assertEqual(ParseUnit('0'), 0)

    self.assertEqual(ParseUnit('0.5'), 4)

    self.assertEqual(ParseUnit('1.75'), 4)

    self.assertEqual(ParseUnit('1.99'), 4)

    self.assertEqual(ParseUnit('2.00'), 4)

    self.assertEqual(ParseUnit('2.01'), 4)

    self.assertEqual(ParseUnit('3.99'), 4)

    self.assertEqual(ParseUnit('4.00'), 4)

    self.assertEqual(ParseUnit('4.01'), 8)

    self.assertEqual(ParseUnit('1.5G'), 1536)

    self.assertEqual(ParseUnit('1.8G'), 1844)

    self.assertEqual(ParseUnit('8.28T'), 8682212)

  def testSuffixes(self):

    for sep in ('', ' ', '   ', "\t", "\t "):

      for suffix, scale in TestParseUnit.SCALES:

        for func in (lambda x: x, str.lower, str.upper):

          self.assertEqual(ParseUnit('1024' + sep + func(suffix)),

                           1024 * scale)

  def testInvalidInput(self):

    for sep in ('-', '_', ',', 'a'):

      for suffix, _ in TestParseUnit.SCALES:

        self.assertRaises(UnitParseError, ParseUnit, '1' + sep + suffix)

    for suffix, _ in TestParseUnit.SCALES:

      self.assertRaises(UnitParseError, ParseUnit, '1,3' + suffix)

class TestSshKeys(testutils.GanetiTestCase):

  """Test case for the AddAuthorizedKey function"""

  KEY_A = 'ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a'

  KEY_B = ('command="/usr/bin/fooserver -t --verbose",from="1.2.3.4" '

           'ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b')

  def setUp(self):

    testutils.GanetiTestCase.setUp(self)

    self.tmpname = self._CreateTempFile()

    handle = open(self.tmpname, 'w')

    try:

      handle.write("%s\n" % TestSshKeys.KEY_A)

      handle.write("%s\n" % TestSshKeys.KEY_B)

    finally:

      handle.close()

  def testAddingNewKey(self):

    AddAuthorizedKey(self.tmpname, 'ssh-dss AAAAB3NzaC1kc3MAAACB root@test')

    self.assertFileContent(self.tmpname,

      "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n"

      'command="/usr/bin/fooserver -t --verbose",from="1.2.3.4"'

      " ssh-dss AAAAB3NzaC1w520smc01ms0jfJs22 root@key-b\n"

      "ssh-dss AAAAB3NzaC1kc3MAAACB root@test\n")

  def testAddingAlmostButNotCompletelyTheSameKey(self):

    AddAuthorizedKey(self.tmpname,

        'ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@test')

    self.assertFileContent(self.tmpname,

      "ssh-dss AAAAB3NzaC1w5256closdj32mZaQU root@key-a\n"