#!/usr/bin/env python

# Copyright (c) 2009, Andrew McNabb
# Copyright (c) 2003-2008, Brent N. Chun

import os
import subprocess
import sys
import shutil
import tempfile
import time
import unittest

basedir, bin = os.path.split(os.path.dirname(os.path.abspath(sys.argv[0])))
sys.path.append("%s" % basedir)

if os.getenv("TEST_HOSTS") is None:
    raise Exception("Must define TEST_HOSTS")
g_hosts = os.getenv("TEST_HOSTS").split()

if os.getenv("TEST_USER") is None:
    raise Exception("Must define TEST_USER")
g_user = os.getenv("TEST_USER")

class PsshTest(unittest.TestCase):
    def setUp(self):
        self.outDir = tempfile.mkdtemp()
        self.errDir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.errDir)
        shutil.rmtree(self.outDir)

    def testShortOpts(self):
        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pssh -h %s -l %s -p 64 -o %s -e %s -t 60 -v -P -i uptime < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        for host in g_hosts:
            stdout = open("%s/%s" % (self.outDir, host)).read()
            self.assert_(stdout.find("load average") != -1)

    def testLongOpts(self):
        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pssh --hosts=%s --user=%s --par=64 --outdir=%s --errdir=%s --timeout=60 --verbose --print --inline uptime < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        for host in g_hosts:
            stdout = open("%s/%s" % (self.outDir, host)).read()
            self.assert_(stdout.find("load average") != -1)

    def testStderr(self):
        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pssh -h %s -l %s -p 64 -o %s -e %s -t 60 -v -P -i ls /doesnotexist < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 5)
        for host in g_hosts:
            stderr = open("%s/%s" % (self.errDir, host)).read()
            self.assert_(stderr.find("No such file or directory") != -1)

class PscpTest(unittest.TestCase):
    def setUp(self):
        self.outDir = tempfile.mkdtemp()
        self.errDir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.errDir)
        shutil.rmtree(self.outDir)
        try:
            os.remove("/tmp/pssh.test")
        except OSError:
            pass

    def testShortOpts(self):
        for host in g_hosts:
            cmd = "ssh %s@%s rm -rf /tmp/pssh.test" % (g_user, host)
            rv = subprocess.call(cmd, shell=True)
            self.assertEqual(rv, 0)

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pscp -h %s -l %s -p 64 -o %s -e %s -t 60 /etc/hosts /tmp/pssh.test < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        for host in g_hosts:
            cmd = "ssh %s@%s cat /tmp/pssh.test" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate()
            self.assertEqual(out, open("/etc/hosts").read())

    def testLongOpts(self):
        for host in g_hosts:
            cmd = "ssh %s@%s rm -rf /tmp/pssh.test" % (g_user, host)
            rv = subprocess.call(cmd, shell=True)
            self.assertEqual(rv, 0)

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pscp --hosts=%s --user=%s --par=64 --outdir=%s --errdir=%s --timeout=60 /etc/hosts /tmp/pssh.test < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        for host in g_hosts:
            cmd = "ssh %s@%s cat /tmp/pssh.test" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate()
            self.assertEqual(out, open("/etc/hosts").read())

    def testRecursive(self):
        for host in g_hosts:
            cmd = "ssh %s@%s rm -rf /tmp/pssh.test" % (g_user, host)
            rv = subprocess.call(cmd, shell=True)
            self.assertEqual(rv, 0)

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pscp -r -h %s -l %s -p 64 -o %s -e %s -t 60 %s/test/test_dir /tmp/pssh.test < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir, basedir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        proc = subprocess.Popen("ls -R %s/test/test_dir | sed -e 's|%s/test/test_dir||' | sort" % (basedir, basedir), stdout=subprocess.PIPE, shell=True)
        (files, err) = proc.communicate()
        for host in g_hosts:
            cmd = "ssh %s@%s ls -R /tmp/pssh.test | sed -e 's|/tmp/pssh.test||' | sort" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate()
            self.assertEqual(out.strip(), files.strip())

class PslurpTest(unittest.TestCase):
    def setUp(self):
        self.outDir = tempfile.mkdtemp()
        self.errDir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.errDir)
        shutil.rmtree(self.outDir)

    def testShortOpts(self):
        if os.path.exists("/tmp/pssh.test"):
            try:
                os.remove("/tmp/pssh.test")
            except OSError:
                shutil.rmtree("/tmp/pssh.test")

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pslurp -L /tmp/pssh.test -h %s -l %s -p 64 -o %s -e %s -t 60 /etc/hosts hosts < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)

        for host in g_hosts:
            cmd = "ssh %s@%s cat /etc/hosts" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate()
            self.assertEqual(out, open("/tmp/pssh.test/%s/hosts" % host).read())

    def testLongOpts(self):
        if os.path.exists("/tmp/pssh.test"):
            try:
                os.remove("/tmp/pssh.test")
            except OSError:
                shutil.rmtree("/tmp/pssh.test")

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pslurp --localdir=/tmp/pssh.test --hosts=%s --user=%s --par=64 --outdir=%s --errdir=%s --timeout=60 /etc/hosts hosts < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)

        for host in g_hosts:
            cmd = "ssh %s@%s cat /etc/hosts" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate()
            self.assertEqual(out, open("/tmp/pssh.test/%s/hosts" % host).read())

    def testRecursive(self):
        if os.path.exists("/tmp/pssh.test"):
            try:
                os.remove("/tmp/pssh.test")
            except OSError:
                shutil.rmtree("/tmp/pssh.test")

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pslurp -r -L /tmp/pssh.test -h %s -l %s -p 64 -o %s -e %s -t 60 /etc/init.d init.d < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)

        for host in g_hosts:
            cmd = "ssh %s@%s ls -R /etc/init.d | sed 1d | sort" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (remoteout, err) = proc.communicate(cmd)

            cmd = "ls -R /tmp/pssh.test/%s/init.d | sed 1d | sort" % host
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (localout, err) = proc.communicate(cmd)

            self.assertEqual(remoteout, localout)

class PrsyncTest(unittest.TestCase):
    def setUp(self):
        self.outDir = tempfile.mkdtemp()
        self.errDir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.errDir)
        shutil.rmtree(self.outDir)

    def testShortOpts(self):
        for host in g_hosts:
            cmd = "ssh %s@%s rm -rf /tmp/pssh.test" % (g_user, host)
            rv = subprocess.call(cmd, shell=True)
            self.assertEqual(rv, 0)

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/prsync -h %s -l %s -p 64 -o %s -e %s -t 60 -a -z /etc/hosts /tmp/pssh.test < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        for host in g_hosts:
            cmd = "ssh %s@%s cat /tmp/pssh.test" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate(cmd)
            self.assertEqual(out, open("/etc/hosts").read())

    def testLongOpts(self):
        for host in g_hosts:
            cmd = "ssh %s@%s rm -rf /tmp/pssh.test" % (g_user, host)
            rv = subprocess.call(cmd, shell=True)
            self.assertEqual(rv, 0)

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/prsync --hosts=%s --user=%s --par=64 --outdir=%s --errdir=%s --timeout=60 --archive --compress /etc/hosts /tmp/pssh.test < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        for host in g_hosts:
            cmd = "ssh %s@%s cat /tmp/pssh.test" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (out, err) = proc.communicate(cmd)
            self.assertEqual(out, open("/etc/hosts").read())

    def testRecursive(self):
        for host in g_hosts:
            cmd = "ssh %s@%s rm -rf /tmp/pssh.test" % (g_user, host)
            rv = subprocess.call(cmd, shell=True)
            self.assertEqual(rv, 0)

        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/prsync -r -h %s -l %s -p 64 -o %s -e %s -t 60 -a -z %s/test/test_dir/ /tmp/pssh.test/ < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir, basedir)
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)
        cmd = "ls -R %s/test/test_dir | sed -e 's|%s/test/test_dir||' | sort" % (basedir, basedir)
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
        (localout, err) = proc.communicate(cmd)
        for host in g_hosts:
            cmd = "ssh %s@%s ls -R /tmp/pssh.test | sed -e 's|/tmp/pssh.test||' | sort" % (g_user, host)
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            (remoteout, err) = proc.communicate(cmd)
            self.assertEqual(localout, remoteout)

class PnukeTest(unittest.TestCase):
    def setUp(self):
        self.outDir = tempfile.mkdtemp()
        self.errDir = tempfile.mkdtemp()

    def tearDown(self):
        shutil.rmtree(self.errDir)
        shutil.rmtree(self.outDir)

    def testShortOpts(self):
        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pssh -h %s -l %s -p 64 -o %s -e %s -t 60 -v sleep 60 < /dev/null &" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        subprocess.call(cmd, shell=True)
        time.sleep(5)

        cmd = "%s/bin/pnuke -h %s -l %s -p 64 -o %s -e %s -t 60 -v sleep < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        print cmd
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)

    def testLongOpts(self):
        hostsFile = tempfile.NamedTemporaryFile()
        hostsFile.write("".join(map(lambda x: "%s\n" % x, g_hosts)))
        hostsFile.flush()
        cmd = "%s/bin/pssh --hosts=%s --user=%s --par=64 --outdir=%s --errdir=%s --timeout=60 --verbose sleep 60 < /dev/null &" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        subprocess.call(cmd, shell=True)
        time.sleep(5)

        cmd = "%s/bin/pnuke --hosts=%s --user=%s --par=64 --outdir=%s --errdir=%s --timeout=60 --verbose sleep < /dev/null" % (basedir, hostsFile.name, g_user, self.outDir, self.errDir)
        print cmd
        rv = subprocess.call(cmd, shell=True)
        self.assertEqual(rv, 0)

if __name__ == '__main__':
    suite = unittest.TestSuite()
    suite.addTest(unittest.makeSuite(PsshTest, "test"))
    suite.addTest(unittest.makeSuite(PscpTest, "test"))
    suite.addTest(unittest.makeSuite(PslurpTest, "test"))
    suite.addTest(unittest.makeSuite(PrsyncTest, "test"))
    suite.addTest(unittest.makeSuite(PnukeTest, "test"))
    unittest.TextTestRunner().run(suite)
