#!/usr/bin/env python
#
# $Id$
#

"""
Note: this is targeted for python 2.x.
To run it under python 3.x you need to use 2to3 tool first:

$ 2to3 -w test/test_memory_leaks.py
"""


import os
import gc
import unittest
import time

import psutil
from test_psutil import reap_children, skipUnless, skipIf, \
                        POSIX, LINUX, WINDOWS, OSX, BSD, PY3

LOOPS = 1000
TOLERANCE = 4096

if PY3:
    xrange = range


class Base(unittest.TestCase):

    def execute(self, function, *args, **kwargs):
        # step 1
        for x in xrange(LOOPS):
            self.call(function, *args, **kwargs)
        del x; gc.collect()
        rss1 = self.get_mem()

        # step 2
        for x in xrange(LOOPS):
            self.call(function, *args, **kwargs)
        del x; gc.collect()
        rss2 = self.get_mem()

        # comparison
        difference = rss2 - rss1
        if difference > TOLERANCE:
            # This doesn't necessarily mean we have a leak yet.
            # At this point we assume that after having called the
            # function so many times the memory usage is stabilized
            # and if there are no leaks it should not increase any
            # more.
            # Let's keep calling fun for 3 more seconds and fail if
            # we notice any difference.
            stop_at = time.time() + 3
            while 1:
                self.call(function, *args, **kwargs)
                if time.time() >= stop_at:
                    break
            rss3 = self.get_mem()
            difference = rss3 - rss2
            if rss3 > rss2:
                self.fail("rss2=%s, rss3=%s, difference=%s" \
                            % (rss2, rss3, difference))

    @staticmethod
    def get_mem():
        return psutil.Process(os.getpid()).get_memory_info()[0]

    def call(self):
        raise NotImplementedError("must be implemented in subclass")


class TestProcessObjectLeaks(Base):
    """Test leaks of Process class methods and properties"""

    def setUp(self):
        gc.collect()

    def tearDown(self):
        reap_children()

    def call(self, function, *args, **kwargs):
        p = psutil.Process(os.getpid())
        obj = getattr(p, function)
        if callable(obj):
            obj(*args, **kwargs)

    def test_name(self):
        self.execute('name')

    def test_cmdline(self):
        self.execute('cmdline')

    def test_ppid(self):
        self.execute('ppid')

    @skipIf(WINDOWS)
    def test_uids(self):
        self.execute('uids')

    @skipIf(WINDOWS)
    def test_gids(self):
        self.execute('gids')

    def test_status(self):
        self.execute('status')

    @skipIf(POSIX)
    def test_username(self):
        self.execute('username')

    def test_create_time(self):
        self.execute('create_time')

    def test_get_num_threads(self):
        self.execute('get_num_threads')

    def test_get_threads(self):
        self.execute('get_num_threads')

    def test_get_cpu_times(self):
        self.execute('get_cpu_times')

    def test_get_memory_info(self):
        self.execute('get_memory_info')

    def test_is_running(self):
        self.execute('is_running')

    @skipIf(WINDOWS)
    def test_terminal(self):
        self.execute('terminal')

    @skipUnless(WINDOWS)
    def test_resume(self):
        self.execute('resume')

    @skipUnless(WINDOWS)
    def test_getcwd(self):
        self.execute('getcwd')

    @skipUnless(WINDOWS or OSX)
    def test_get_open_files(self):
        self.execute('get_open_files')

    @skipUnless(WINDOWS or OSX)
    def test_get_connections(self):
        self.execute('get_connections')


class TestModuleFunctionsLeaks(Base):
    """Test leaks of psutil module functions."""

    def setUp(self):
        gc.collect()

    def call(self, function, *args, **kwargs):
        obj = getattr(psutil, function)
        if callable(obj):
            retvalue = obj(*args, **kwargs)

    def test_get_pid_list(self):
        self.execute('get_pid_list')

    @skipIf(POSIX)
    def test_pid_exists(self):
        self.execute('pid_exists', os.getpid())

    def test_process_iter(self):
        self.execute('process_iter')

    def test_phymem_usage(self):
        self.execute('phymem_usage')

    def test_virtmem_usage(self):
        self.execute('virtmem_usage')

    def test_cpu_times(self):
        self.execute('cpu_times')

    def test_per_cpu_times(self):
        self.execute('per_cpu_times')

    @skipUnless(WINDOWS)
    def test_disk_usage(self):
        self.execute('disk_usage', '.')

    def test_disk_partitions(self):
        self.execute('disk_partitions')


def test_main():
    test_suite = unittest.TestSuite()
    test_suite.addTest(unittest.makeSuite(TestProcessObjectLeaks))
    test_suite.addTest(unittest.makeSuite(TestModuleFunctionsLeaks))
    unittest.TextTestRunner(verbosity=2).run(test_suite)

if __name__ == '__main__':
    test_main()


@@ -604,33 +604,40 @@
 s)%0A        del x
-;
+%0A       
  gc.collect()%0A  
@@ -780,9 +780,16 @@
 el x
-;
+%0A       
  gc.
@@ -1493,16 +1493,65 @@
   break%0A
+            del stop_at%0A            gc.collect()%0A
         
@@ -1703,18 +1703,16 @@
 e=%25s%22 %5C%0A
-  
         

call gc.collect() in step 3