mock 示例

import sys,os
# workdir=os.getcwd()
# root_path=os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# print ("root_path",root_path)
# sys.path.append(root_path)

from unittest.mock import Mock
from ..index_manage.screenIndex import IndexManage
from pandas import DataFrame
import pandas as pd
import psycopg2
import psycopg2.extras

class FengClass(object):
    wfids = '0cb03afdc42e443e947506146edc80be,1004'  # feng
    stime = '2020-04-19'
    etime = '2020-04-24'
    cols = ""

class GuangClass(object):
    wfids = '57bacf6dc8b54b7e963385ec0598f6lp,1005'  # guang
    stime = '2020-04-19'
    etime = '2020-04-24'
    cols = ""

class GuangErrorClass(object):
    wfids = '57bacf6dc8b54b7e963385ec0598f6lp,1005'  # guang
    stime = '2020-04-19'
    etime = '2020-04-24'
    cols = "sdf"

class noWfidClass(object):
    wfids = ''  # guang
    stime = '2020-04-19'
    etime = '2020-04-24'
    cols = ""

class myCurs(object):
    def __enter__(self):

        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def execute(self,sql):
        pass
    def fetchall(self):
        t = [11] * 96
        t.insert(0, '2020-01-01')
        return [t]
    def fetchone(self):
        return "3eee6536043e44ecbadb525a6e072e95"

class AccessDB(object):
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def close(self):
        pass
    def cursor(self):
        return myCurs()

class Test_tool():
    def setup_method(self): # 一开始用的setup_class，结果一个用例的mock影响了下个用例，因为这些用例共用了一个self.tstobj。现在改成setup_method后，每个用例执行前都重置self.tstobj的值

        self.tstobj = IndexManage()

    def test_invalid_col(self):
        psycopg2.connect = Mock(return_value=None)
        self.tstobj.write_db = Mock(return_value=None)
        self.tstobj.get_neibs = Mock(return_value=['3eee6536043e44ecbadb525a6e072e95'])
        data_x = DataFrame({'x': [1, 2, 3, 4]}, index = pd.date_range('2020-02-04', periods=4, freq="15T"))
        data_y = DataFrame({'y': [2, 2, 4, 4]}, index = pd.date_range('2020-02-04', periods=4, freq="15T"))
        self.tstobj.get_value = Mock(side_effect=lambda _, x, *d: data_x if 'x' == x else data_y)
        ret = self.tstobj.corr(GuangErrorClass, None)
        assert "" == ret.ret


    def test_execute_batch_error(self):
        psycopg2.connect = Mock(return_value=AccessDB())
        self.tstobj.get_neibs = Mock(return_value=['3eee6536043e44ecbadb525a6e072e95'])
        data_x = DataFrame({'x': [1, 2, 3, 4]}, index=pd.date_range('2020-02-04', periods=4, freq="15T"))
        data_y = DataFrame({'y': [2, 2, 4, 4]}, index=pd.date_range('2020-02-04', periods=4, freq="15T"))
        self.tstobj.get_value = Mock(side_effect=lambda _, x, *d: data_x if 'x' == x else data_y)
        ret = self.tstobj.corr(GuangClass, None)
        assert "" == ret.ret


    def test_get_value(self):
        psycopg2.connect = Mock(return_value=AccessDB())
        psycopg2.extras.execute_batch = Mock(return_value = None)
        ret = self.tstobj.corr(GuangClass, None)
        assert "" == ret.ret

    def test_get_value_xnone(self):
        psycopg2.connect = Mock(return_value=None)
        self.tstobj.write_db = Mock(return_value=None)
        self.tstobj.get_neibs = Mock(return_value = ['3eee6536043e44ecbadb525a6e072e95'])
        data_x = DataFrame()
        data_y = DataFrame({'y': [2, 2, 4, 4]}, index=pd.date_range('2020-02-04', periods=4, freq="15T"))
        self.tstobj.get_value = Mock(side_effect=lambda _, x, *d: data_x if 'x' == x else data_y)  # Mock(side_effect=lambda x: 0 if x<20 else 30)
        ret = self.tstobj.corr(FengClass, None)
        assert "" == ret.ret

    def test_get_value_ynone(self):
        psycopg2.connect = Mock(return_value=None)
        self.tstobj.write_db = Mock(return_value=None)
        self.tstobj.get_neibs = Mock(return_value = ['3eee6536043e44ecbadb525a6e072e95'])
        data_x = DataFrame({'x': [1, 2, 3, 4]}, index=pd.date_range('2020-02-04', periods=4, freq="15T"))
        data_y = DataFrame()
        self.tstobj.get_value = Mock(side_effect=lambda _, x, *d: data_x if 'x' == x else data_y)  # Mock(side_effect=lambda x: 0 if x<20 else 30)
        ret = self.tstobj.corr(FengClass, None)
        assert "" == ret.ret


    def test_noWfidClass(self):
        psycopg2.connect = Mock(return_value=AccessDB())
        psycopg2.extras.execute_batch = Mock(return_value=None)
        ret = self.tstobj.corr(noWfidClass, None)

        assert 'unexpect type' in ret.ret



    def test_feng_no_neibs(self):
        psycopg2.connect = Mock(return_value=None)
        self.tstobj.write_db = Mock(return_value=None)
        self.tstobj.get_neibs = Mock(return_value = None)
        ret = self.tstobj.corr(FengClass, None)
        assert "no neighbor" in ret.ret

    def test_feng(self):
        psycopg2.connect = Mock(return_value=None)
        self.tstobj.write_db = Mock(return_value=None)
        self.tstobj.get_neibs = Mock(return_value = ['3eee6536043e44ecbadb525a6e072e95'])
        data_x = DataFrame({'x': [1, 2, 3, 4]}, index=pd.date_range('2020-02-04', periods=4, freq="15T"))
        data_y = DataFrame({'y': [2, 2, 4, 4]}, index=pd.date_range('2020-02-04', periods=4, freq="15T"))
        self.tstobj.get_value = Mock(side_effect=lambda _, x, *d: data_x if 'x' == x else data_y)  # Mock(side_effect=lambda x: 0 if x<20 else 30)
        ret = self.tstobj.corr(FengClass, None)
        assert "" == ret.ret

    def test_guang(self):
        psycopg2.connect = Mock(return_value=None)
        self.tstobj.write_db = Mock(return_value=None)
        self.tstobj.get_neibs = Mock(return_value=['3eee6536043e44ecbadb525a6e072e95'])
        data_x = DataFrame({'x': [1, 2, 3, 4]}, index = pd.date_range('2020-02-04', periods=4, freq="15T"))
        data_y = DataFrame({'y': [2, 2, 4, 4]}, index = pd.date_range('2020-02-04', periods=4, freq="15T"))
        self.tstobj.get_value = Mock(side_effect=lambda _, x, *d: data_x if 'x' == x else data_y)
        ret = self.tstobj.corr(GuangClass, None)
        assert "" == ret.ret