~ubuntu-core-dev/update-notifier/ubuntu

#!/usr/bin/python3

import hashlib

import http.server as SimpleHTTPServer

import os

import random

import shutil

import socketserver

import sys

import time

import tempfile

import unittest

from mock import patch

from multiprocessing import Process

sys.path.insert(0, "data")

sys.path.insert(1, "../data")

import package_data_downloader  # nopep8

class TestHttpd(Process):

    def __init__(self, port, servdir):

        super(TestHttpd, self).__init__()

        self.port = port

        self.servdir = servdir

    def run(self):

        os.chdir(self.servdir)

        Handler = SimpleHTTPServer.SimpleHTTPRequestHandler

        httpd = socketserver.TCPServer(("localhost", self.port), Handler)

        httpd.serve_forever()

class DownloadTestCase(unittest.TestCase):

    def setUp(self):

        self.tempdir = tempfile.mkdtemp()

        self._setup_canary_file()

        self.PORT = random.randint(1025, 60000)

        self.p = TestHttpd(self.PORT, self.tempdir)

        self.p.start()

    def tearDown(self):

        self.p.terminate()

        shutil.rmtree(self.tempdir)

    def _setup_canary_file(self):

        self.canary_text = "meep"

        self.canary_file = os.path.join(self.tempdir, "canary-file.txt")

        with open(self.canary_file, "w") as f:

            f.write(self.canary_text)

        self.canary_hashsum = hashlib.sha256(

            self.canary_text.encode('utf-8')).hexdigest()

    def test_download_file(self):

        package_data_downloader.STAMPDIR = self.tempdir

        os.makedirs(os.path.join(self.tempdir, "partial"))

        uri = "http://localhost:%s/%s" % (

            self.PORT, os.path.basename(self.canary_file))

        destfile = package_data_downloader.download_file(

            uri, self.canary_hashsum)

        self.assertNotEqual(destfile, None)

        self.assertEqual(hashlib.sha256(

            open(destfile).read().encode('utf-8')).hexdigest(),

            self.canary_hashsum)

class ProcessDownloadRequestsTestCase(unittest.TestCase):

    def setUp(self):

        self.tempdir = tempfile.mkdtemp()

        self._setup_canary_file()

        self.PORT = random.randint(1025, 60000)

        self.p = TestHttpd(self.PORT, self.tempdir)

        self.p.start()

        self.tempdir = tempfile.mkdtemp()

        # stampdir

        stampdir = os.path.join(self.tempdir, "stampdir")

        os.makedirs(stampdir)

        package_data_downloader.STAMPDIR = stampdir

        os.makedirs(os.path.join(stampdir, "partial"))

        # datadir

        datadir = os.path.join(self.tempdir, "datadir")

        os.makedirs(datadir)

        package_data_downloader.DATADIR = datadir

        # override the location of the notifier files

        notifierdir = os.path.join(self.tempdir, "notifierdir")

        os.makedirs(notifierdir)

        package_data_downloader.NOTIFIER_FILE = \

            os.path.join(notifierdir, "data-downloads-failed")

        package_data_downloader.NOTIFIER_PERMANENT_FILE = \

            package_data_downloader.NOTIFIER_FILE + "-permanently"

    def tearDown(self):

        self.p.terminate()

        shutil.rmtree(self.tempdir)

    def _setup_canary_file(self):

        self.canary_text = "meep"

        self.canary_file = os.path.join(self.tempdir, "canary-file.txt")

        with open(self.canary_file, "w") as f:

            f.write(self.canary_text)

        self.canary_hashsum = hashlib.sha256(

            self.canary_text.encode('utf-8')).hexdigest()

    def _setup_hook_file(self, filename, script="/bin/true"):

        with open(os.path.join(package_data_downloader.DATADIR,

                  filename), "w") as foo:

            foo.write("Url: http://localhost:%s/%s\n" %

                      (self.PORT, os.path.basename(self.canary_file)))

            foo.write("Sha256: %s\n" % self.canary_hashsum)

            foo.write("\n")

            foo.write("Script: %s" % script)

    def test_hookfile_older_than_stampfile(self):

        # hook file older than stamp file nothing happens

        # create the hook file

        hookfile = "older-hookfile"

        self._setup_hook_file(hookfile)

        time.sleep(0.01)

        # create the stamp file

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        orig_stamp_time = os.stat(stampfile).st_mtime

        package_data_downloader.process_download_requests()

        new_stamp_time = os.stat(stampfile).st_mtime

        self.assertEqual(orig_stamp_time, new_stamp_time)

        # confirm failure files not created

        self.assertFalse(os.path.exists(

            os.path.join(package_data_downloader.STAMPDIR,

                         "%s.permanent-failure" % hookfile)))

        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

    def test_stampfile_older_than_hookfile(self):

        # hook file newer than stampfile, download, run script, update

        # stampfile

        hookfile = "older-stampfile"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        orig_stamp_date = os.stat(stampfile).st_mtime

        time.sleep(0.01)

        # create the hook file

        self._setup_hook_file(hookfile)

        package_data_downloader.process_download_requests()

        # the download succeeded so the stamp file is touched

        new_stamp_date = os.stat(stampfile).st_mtime

        self.assertGreater(new_stamp_date, orig_stamp_date)

        # confirm failure files not created

        self.assertFalse(os.path.exists(

            os.path.join(package_data_downloader.STAMPDIR,

                         "%s.permanent-failure" % hookfile)))

        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

    def test_only_retry_failure(self):

        # two hook files but only has failed

        hookfile = "success"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

            pass

        orig_stamp_date = os.stat(stampfile).st_mtime

        time.sleep(0.01)

        # create the hook file

        self._setup_hook_file(hookfile)

        fail_hookfile = "failure"

        # the stampfile indicates a failure

        fail_stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                      "%s.failed" % fail_hookfile)

        with open(fail_stampfile, "w"):

            pass

        time.sleep(0.01)

        # create the hook file

        self._setup_hook_file(fail_hookfile)

        # create an empty notifier file

        with open(package_data_downloader.NOTIFIER_FILE, "w"):

            pass

        package_data_downloader.process_download_requests()

        new_stamp_date = os.stat(stampfile).st_mtime

        # if the downloaded succeeded it shouldn't be done again

        self.assertEqual(new_stamp_date, orig_stamp_date)

    def test_hookfile_script_fails(self):

        # script in the hook file fails, failure recorded as permanent

        hookfile = "script-failure"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        time.sleep(0.01)

        # create the hook file

        self._setup_hook_file(hookfile, "/bin/false")

        package_data_downloader.process_download_requests()

        # script failures are considered permanent

        self.assertTrue(os.path.exists(

            os.path.join(package_data_downloader.STAMPDIR,

                         "%s.permanent-failure" % hookfile)))

    def test_stampfile_and_notifierfile(self):

        # notifier file removed on successful download

        hookfile = "stampfile_and_notifierfile"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        time.sleep(0.01)

        self._setup_hook_file(hookfile)

        # create an empty notifier file

        with open(package_data_downloader.NOTIFIER_FILE, "w"):

            pass

        package_data_downloader.process_download_requests()

        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

    def test_stampfile_notifierfile_and_notifierpermanent(self):

        # both notifier files removed on successful download

        hookfile = "stampfile_notifierfile_and_notifierpermanentfile"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        time.sleep(0.01)

        self._setup_hook_file(hookfile)

        # create empty notifier files

        with open(package_data_downloader.NOTIFIER_FILE, "w"):

            pass

        with open(package_data_downloader.NOTIFIER_PERMANENT_FILE, "w"):

            pass

        package_data_downloader.process_download_requests()

        self.assertFalse(os.path.exists(package_data_downloader.NOTIFIER_FILE))

        self.assertFalse(os.path.exists(

            package_data_downloader.NOTIFIER_PERMANENT_FILE))

    def test_stampfile_and_notifierpermanent(self):

        # permanent notifier file removed on successful download

        hookfile = "stampfile_and_permanentnotifierfile"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        time.sleep(0.01)

        self._setup_hook_file(hookfile)

        # create an empty permanent notifier file

        with open(package_data_downloader.NOTIFIER_PERMANENT_FILE, "w"):

            pass

        package_data_downloader.process_download_requests()

        self.assertFalse(os.path.exists(

            package_data_downloader.NOTIFIER_PERMANENT_FILE))

    def test_hookfile_download_failure(self):

        # can't download the file, non-permanent failure created

        hookfile = "download-failure"

        stampfile = os.path.join(package_data_downloader.STAMPDIR,

                                 hookfile)

        with open(stampfile, "w"):

                pass

        time.sleep(0.01)

        # overwrite canary file to create a failure

        self.canary_file = "not-there.txt"

        # create the hook file

        self._setup_hook_file(hookfile)

        package_data_downloader.process_download_requests()

        self.assertTrue(os.path.exists(

            os.path.join(package_data_downloader.STAMPDIR,

                         "%s.failed" % hookfile)))

        self.assertFalse(os.path.exists(

            os.path.join(package_data_downloader.STAMPDIR,

                         "%s.permanent-failure" % hookfile)))

    def test_hookfile_notifierfile_download_failure(self):

        # can't download the file, notifier file stays around

        hookfile = "download-failure-with-notifierfile"

        # overwrite canary file to create a failure

        self.canary_file = "not-there.txt"

        # create the hook file

        self._setup_hook_file(hookfile)

        # create an empty notifier file

        with open(package_data_downloader.NOTIFIER_FILE, "w"):

            pass

        package_data_downloader.process_download_requests()

        # because there was a failure it shouldn't be removed

        self.assertTrue(os.path.exists(

            package_data_downloader.NOTIFIER_FILE))

class PackageDataDownloaderTestCase(unittest.TestCase):

    def setUp(self):

        self.tmpdir = tempfile.mkdtemp()

        # stampdir

        stampdir = os.path.join(self.tmpdir, "stampdir")

        os.makedirs(stampdir)

        package_data_downloader.STAMPDIR = stampdir

        # datadir

        datadir = os.path.join(self.tmpdir, "datadir")

        os.makedirs(datadir)

        package_data_downloader.DATADIR = datadir

    def tearDown(self):

        shutil.rmtree(self.tmpdir)

    def test_wrong_template_translations(self):

        package_data_downloader.NOTIFIER_SOURCE_FILE = \

            'data/package-data-downloads-failed.in'

        package_data_downloader.NOTIFIER_FILE = \

            self.tmpdir + "/data-downloads-failed"

        package_data_downloader.NOTIFIER_PERMANENT_SOURCE_FILE = \

            'data/package-data-downloads-failed-permanently.in'

        package_data_downloader.NOTIFIER_PERMANENT_FILE = \

            package_data_downloader.NOTIFIER_FILE + '-permanently'

        package_data_downloader.trigger_update_notifier([], False)

        package_data_downloader.trigger_update_notifier([], True)

    def test_permanently_failed(self):

        # create a bunch of files using the provided mechanism

        test_files = ["foo.permanent-failure", "bar.failure", "baz"]

        for f in test_files:

            package_data_downloader.create_or_update_stampfile(

                os.path.join(package_data_downloader.STAMPDIR, f))

        self.assertEqual(

            sorted(os.listdir(package_data_downloader.STAMPDIR)),

            sorted(test_files))

        # test hook_is_permanently_failed()

        self.assertTrue(

            package_data_downloader.hook_is_permanently_failed("foo"))

        self.assertFalse(

            package_data_downloader.hook_is_permanently_failed("bar"))

        self.assertFalse(

            package_data_downloader.hook_is_permanently_failed("baz"))

        # existing_permanent_failures()

        self.assertEqual(

            package_data_downloader.existing_permanent_failures(),

            ["foo"])

    def test_hook_aged_out(self):

        # test to see if a hook has failed for more than 3 days

        # create a failure file

        with open(os.path.join(package_data_downloader.STAMPDIR,

                               "aged.failed"), "w"):

            pass

        from datetime import datetime, timedelta

        # patch datetime so we think its the future

        with patch('package_data_downloader.datetime') as mock_datetime:

            thefuture = (datetime.now() + timedelta(days=3))

            mock_datetime.now.return_value = thefuture

            mock_datetime.fromtimestamp.side_effect = \

                lambda *args, **kw: datetime.fromtimestamp(*args, **kw)

            self.assertEqual(package_data_downloader.datetime.now(),

                             thefuture)

            self.assertTrue(package_data_downloader.hook_aged_out("aged"))

    def test_mark_hook_failed(self):

        # prepare

        package_data_downloader.create_or_update_stampfile(

            os.path.join(package_data_downloader.STAMPDIR, "foo"))

        # temp failure

        package_data_downloader.mark_hook_failed("foo")

        self.assertEqual(os.listdir(package_data_downloader.STAMPDIR),

                         ["foo.failed"])

        self.assertFalse(

            package_data_downloader.hook_is_permanently_failed("foo"))

        # permanent

        package_data_downloader.mark_hook_failed("foo", permanent=True)

        self.assertEqual(os.listdir(package_data_downloader.STAMPDIR),

                         ["foo.permanent-failure"])

        self.assertTrue(

            package_data_downloader.hook_is_permanently_failed("foo"))

    def test_get_hook_file_names(self):

        # test that .dpkg-* is ignored

        test_datadir_files = ["foo.dpkg-new", "bar"]

        for name in test_datadir_files:

            with open(os.path.join(package_data_downloader.DATADIR, name),

                      "w"):

                pass

        self.assertEqual(package_data_downloader.get_hook_file_names(),

                         ["bar"])

    def test_trigger_update_notifier(self):

        # write to tmpdir

        package_data_downloader.NOTIFIER_FILE = os.path.join(

            self.tmpdir, "data-downloads-failed")

        # point to local repo file

        package_data_downloader.NOTIFIER_SOURCE_FILE = \

            "data/package-data-downloads-failed.in"

        package_data_downloader.trigger_update_notifier(

            failures=["foo", "bar"], permanent=False)

        data = open(package_data_downloader.NOTIFIER_FILE).read()

        self.assertEqual(data, """Priority: High

_Name: Failure to download extra data files

Terminal: True

Command: pkexec /usr/lib/update-notifier/package-data-downloader

_Description:

 The following packages requested additional data downloads after package

 installation, but the data could not be downloaded or could not be

 processed.

 .

 .

   foo, bar

 .

 .

 The download will be attempted again later, or you can try the download

 again now.  Running this command requires an active Internet connection.

""")

if __name__ == "__main__":

    unittest.main()