import imp
import os
import pprint
import sys
from test import test_support 
import unittest

import trace
from trace import CoverageResults, Trace


#------------------------------- Utilities -----------------------------------#

def make_fake_module():
    """Creates a fake module named 'fakemodule'.

    The new module has a single function named 'func', and it's placed in 
    sys.modules
    The file this fake module "comes from" is fakefile.py
    
    """
    
    # Prepare the function to import from the fake module
    # 
    fake_func_src = r'''
def func(a_):
    b = a_ + 1
    return b + 2
'''.lstrip()

    fake_func = compile(fake_func_src, 'fakefile.py', 'exec')
    
    # Create a new module, place the function into it and add it to sys.modules
    #
    fakemodule = imp.new_module('fakemodule')
    exec fake_func in fakemodule.__dict__
    fakemodule.__file__ = 'fakefile.py'
    sys.modules['fakemodule'] = fakemodule


def modname(filename):
    """Infer a module name from a containing file name"""
    base = os.path.basename(filename)
    mod, ext = os.path.splitext(base)
    return mod


def my_file_and_modname():
    """The file and module name of this file (__file__)"""
    return __file__, modname(__file__)


#-------------------- Target functions for tracing ---------------------------#

def traced_func_linear(aa, bb):
    a = aa
    b = bb
    c = a + b
    return c

def traced_func_loop(aa, bb):
    c = aa
    for i in range(5):
        c += bb
    return c

# Expects the 'fakemodule' module to exist and have a 'func' function in it
#
def traced_func_importing(aa, bb):
    from fakemodule import func
    return aa + bb + func(1)

def traced_func_simple_caller(aa):
    c = traced_func_linear(aa, aa)
    return c + aa

def traced_func_importing_caller(aa):
    k = traced_func_simple_caller(aa)
    k += traced_func_importing(k, aa)
    return k

def traced_func_generator(num):
    c = 5       # executed once
    for i in range(num):
        yield i + c

def traced_func_calling_generator():
    k = 0
    for i in traced_func_generator(10):
        k += i

def traced_doubler(num):
    return num * 2

def traced_caller_list_comprehension():
    k = 10
    mylist = [traced_doubler(i) for i in range(k)]
    return mylist


class TracedClass(object):
    def __init__(self, aa):
        self.a = aa
    
    def inst_method_linear(self, bb):
        return self.a + bb
    
    def inst_method_calling(self, aa):
        c = self.inst_method_linear(aa)
        return c + traced_func_linear(aa, c)
    
    @classmethod
    def class_method_linear(cls, bb):
        return bb * 2
    
    @staticmethod
    def static_method_linear(bb):
        return bb * 2


#------------------------------ Test cases -----------------------------------#


class TestLineCounts(unittest.TestCase):
    """White-box testing of line-counting, via runfunc"""
    def setUp(self):
        self.tracer = Trace(count=1, trace=0, countfuncs=0, countcallers=0)

    def test_traced_func_linear(self):
        result = self.tracer.runfunc(traced_func_linear, 2, 5)        
        self.assertEqual(result, 7)
        
        # all lines are executed once
        expected = {}
        firstlineno = traced_func_linear.__code__.co_firstlineno
        for i in range(1, 5):
            expected[(__file__, firstlineno +  i)] = 1
        
        self.assertEqual(self.tracer.results().counts, expected)
        
    def test_traced_func_loop(self):
        self.tracer.runfunc(traced_func_loop, 2, 3)
        
        firstlineno = traced_func_loop.__code__.co_firstlineno
        expected = {
            (__file__, firstlineno + 1): 1,
            (__file__, firstlineno + 2): 6,
            (__file__, firstlineno + 3): 5,
            (__file__, firstlineno + 4): 1,
        }
        self.assertEqual(self.tracer.results().counts, expected)
    
    def test_traced_func_importing(self):
        make_fake_module()
        self.tracer.runfunc(traced_func_importing, 2, 5)        
        
        firstlineno = traced_func_importing.__code__.co_firstlineno
        expected = {
            (__file__, firstlineno + 1): 1,
            (__file__, firstlineno + 2): 1,
            ('fakefile.py', 2): 1,
            ('fakefile.py', 3): 1,
        }
        self.assertEqual(self.tracer.results().counts, expected)

    def test_trace_func_generator(self):
        self.tracer.runfunc(traced_func_calling_generator)

        firstlineno_calling = traced_func_calling_generator.__code__.co_firstlineno
        firstlineno_gen = traced_func_generator.__code__.co_firstlineno
        expected = {
            (__file__, firstlineno_calling + 1): 1,
            (__file__, firstlineno_calling + 2): 11,
            (__file__, firstlineno_calling + 3): 10,
            (__file__, firstlineno_gen + 1): 1,
            (__file__, firstlineno_gen + 2): 11,
            (__file__, firstlineno_gen + 3): 10,
        }
        self.assertEqual(self.tracer.results().counts, expected)

    def test_trace_list_comprehension(self):
        self.tracer.runfunc(traced_caller_list_comprehension)

        firstlineno_calling = traced_caller_list_comprehension.__code__.co_firstlineno
        firstlineno_called = traced_doubler.__code__.co_firstlineno
        expected = {
            (__file__, firstlineno_calling + 1): 1,
            (__file__, firstlineno_calling + 2): 11,
            (__file__, firstlineno_calling + 3): 1,
            (__file__, firstlineno_called + 1): 10,
        }
        self.assertEqual(self.tracer.results().counts, expected)


    def test_linear_methods(self):
        # ZZZ todo: later add 'static_method_linear' and 'class_method_linear'
        # here, once the line numbers of those are resolved
        #
        for methname in ['inst_method_linear',]:
            tracer = Trace(count=1, trace=0, countfuncs=0, countcallers=0)
            traced_obj = TracedClass(25)
            method = getattr(traced_obj, methname)
            tracer.runfunc(method, 20)
        
            firstlineno = method.__code__.co_firstlineno
            expected = {
                (__file__, firstlineno + 1): 1,
            }
            self.assertEqual(tracer.results().counts, expected)
        

class TestRunExecCounts(unittest.TestCase):
    """A simple sanity test of line-counting, via runctx (exec)"""
    def test_exec_counts(self):
        self.tracer = Trace(count=1, trace=0, countfuncs=0, countcallers=0)
        code = r'''traced_func_loop(2, 5)'''
        code = compile(code, __file__, 'exec')
        self.tracer.runctx(code, globals(), vars())
        
        firstlineno = traced_func_loop.__code__.co_firstlineno
        expected = {
            (__file__, firstlineno + 1): 1,
            (__file__, firstlineno + 2): 6,
            (__file__, firstlineno + 3): 5,
            (__file__, firstlineno + 4): 1,
        }
        
        # When used through 'run', some other spurios counts are produced, like
        # the settrace of threading, which we ignore, just making sure that the
        # counts fo traced_func_loop were right.
        #
        for k in expected.keys():
            self.assertEqual(self.tracer.results().counts[k], expected[k])


class TestFuncs(unittest.TestCase):
    """White-box testing of funcs tracing"""
    def setUp(self):
        self.tracer = Trace(count=0, trace=0, countfuncs=1)
        self.filemod = my_file_and_modname()
        
    def test_simple_caller(self):
        self.tracer.runfunc(traced_func_simple_caller, 1)
        
        expected = {
            self.filemod + ('traced_func_simple_caller',): 1,
            self.filemod + ('traced_func_linear',): 1,
        }
        self.assertEqual(self.tracer.results().calledfuncs, expected)

    def test_loop_caller_importing(self):
        make_fake_module()
        self.tracer.runfunc(traced_func_importing_caller, 1)
        
        expected = { 
            self.filemod + ('traced_func_simple_caller',): 1,
            self.filemod + ('traced_func_linear',): 1,
            self.filemod + ('traced_func_importing_caller',): 1,
            self.filemod + ('traced_func_importing',): 1,
            ('fakefile.py', 'fakefile', 'func'): 1,
        }
        self.assertEqual(self.tracer.results().calledfuncs, expected)

    def test_inst_method_calling(self):
        obj = TracedClass(20)
        self.tracer.runfunc(obj.inst_method_calling, 1)
        
        expected = {
            self.filemod + ('TracedClass.inst_method_calling',): 1,
            self.filemod + ('TracedClass.inst_method_linear',): 1,
            self.filemod + ('traced_func_linear',): 1,
        }
        self.assertEqual(self.tracer.results().calledfuncs, expected)        
    

class TestCallers(unittest.TestCase):
    """White-box testing of callers tracing"""
    def setUp(self):
        self.tracer = Trace(count=0, trace=0, countcallers=1)
        self.filemod = my_file_and_modname()
    
    def test_loop_caller_importing(self):
        make_fake_module()
        self.tracer.runfunc(traced_func_importing_caller, 1)

        expected = {
            ((os.path.splitext(trace.__file__)[0] + '.py', 'trace', 'Trace.runfunc'), 
                (self.filemod + ('traced_func_importing_caller',))): 1,
            ((self.filemod + ('traced_func_simple_caller',)), 
                (self.filemod + ('traced_func_linear',))): 1,
            ((self.filemod + ('traced_func_importing_caller',)), 
                (self.filemod + ('traced_func_simple_caller',))): 1,
            ((self.filemod + ('traced_func_importing_caller',)), 
                (self.filemod + ('traced_func_importing',))): 1,
            ((self.filemod + ('traced_func_importing',)), 
                ('fakefile.py', 'fakefile', 'func')): 1,
        }
        self.assertEqual(self.tracer.results().callers, expected)


#------------------------------ Driver ---------------------------------------#

def test_main():
    test_support.run_unittest(__name__)


if __name__ == '__main__':
    test_main()