github.com/bytedance/mockey@v1.2.10/mock_var.go (about)

     1  /*
     2   * Copyright 2022 ByteDance Inc.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package mockey
    18  
    19  import (
    20  	"reflect"
    21  	"sync"
    22  
    23  	"github.com/bytedance/mockey/internal/tool"
    24  )
    25  
    26  type mockerInstance interface {
    27  	key() uintptr
    28  	name() string
    29  	unPatch()
    30  
    31  	caller() tool.CallerInfo
    32  }
    33  
    34  type MockerVar struct {
    35  	target     reflect.Value // 目标变量地址
    36  	hook       reflect.Value // mock变量
    37  	targetType reflect.Type
    38  	origin     interface{} // 原始值
    39  	lock       sync.Mutex
    40  	isPatched  bool
    41  
    42  	outerCaller tool.CallerInfo
    43  }
    44  
    45  func MockValue(targetPtr interface{}) *MockerVar {
    46  	tool.AssertPtr(targetPtr)
    47  
    48  	return &MockerVar{
    49  		target:     reflect.ValueOf(targetPtr).Elem(),
    50  		origin:     reflect.ValueOf(targetPtr).Elem().Interface(),
    51  		targetType: reflect.TypeOf(targetPtr).Elem(),
    52  	}
    53  }
    54  
    55  func (mocker *MockerVar) To(value interface{}) *MockerVar {
    56  	var v reflect.Type
    57  
    58  	if value == nil {
    59  		mocker.hook = reflect.Zero(mocker.targetType)
    60  		v = mocker.targetType
    61  	} else {
    62  		mocker.hook = reflect.ValueOf(value)
    63  		v = reflect.TypeOf(value)
    64  	}
    65  
    66  	tool.Assert(v.AssignableTo(mocker.targetType), "value type:%s not match target type:%s", v.Name(), mocker.targetType.Name())
    67  	mocker.Patch()
    68  	return mocker
    69  }
    70  
    71  func (mocker *MockerVar) Patch() *MockerVar {
    72  	mocker.lock.Lock()
    73  	defer mocker.lock.Unlock()
    74  
    75  	if !mocker.isPatched {
    76  		mocker.target.Set(mocker.hook)
    77  		mocker.isPatched = true
    78  		addToGlobal(mocker)
    79  
    80  		mocker.outerCaller = tool.OuterCaller()
    81  	}
    82  
    83  	return mocker
    84  }
    85  
    86  func (mocker *MockerVar) UnPatch() *MockerVar {
    87  	mocker.lock.Lock()
    88  	defer mocker.lock.Unlock()
    89  	if mocker.isPatched {
    90  		mocker.isPatched = false
    91  		if mocker.origin == nil {
    92  			mocker.target.Set(reflect.Zero(mocker.targetType))
    93  		} else {
    94  			mocker.target.Set(reflect.ValueOf(mocker.origin))
    95  		}
    96  		removeFromGlobal(mocker)
    97  	}
    98  
    99  	return mocker
   100  }
   101  
   102  func (mocker *MockerVar) key() uintptr {
   103  	return mocker.target.Addr().Pointer()
   104  }
   105  
   106  func (mocker *MockerVar) name() string {
   107  	if mocker.target.Kind() == reflect.String {
   108  		return "<string Value>"
   109  	}
   110  	return mocker.target.String()
   111  }
   112  
   113  func (mocker *MockerVar) unPatch() {
   114  	mocker.UnPatch()
   115  }
   116  
   117  func (mocker *MockerVar) caller() tool.CallerInfo {
   118  	return mocker.outerCaller
   119  }