mock.go

Documentation: github.com/chaos-mesh/chaos-mesh/pkg/mock

     1  // Copyright 2020 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package mock
    15  
    16  import (
    17  	"path"
    18  	"reflect"
    19  	"sync"
    20  
    21  	"github.com/pingcap/failpoint"
    22  )
    23  
    24  // Finalizer represent the function that clean a mock point
    25  type Finalizer func() error
    26  
    27  type mockPoints struct {
    28  	m map[string]interface{}
    29  	l sync.Mutex
    30  }
    31  
    32  func (p *mockPoints) set(fpname string, value interface{}) {
    33  	p.l.Lock()
    34  	defer p.l.Unlock()
    35  
    36  	p.m[fpname] = value
    37  }
    38  
    39  func (p *mockPoints) get(fpname string) interface{} {
    40  	p.l.Lock()
    41  	defer p.l.Unlock()
    42  
    43  	return p.m[fpname]
    44  }
    45  
    46  func (p *mockPoints) clr(fpname string) {
    47  	p.l.Lock()
    48  	defer p.l.Unlock()
    49  
    50  	delete(p.m, fpname)
    51  }
    52  
    53  var points = mockPoints{m: make(map[string]interface{})}
    54  
    55  // On inject a failpoint
    56  func On(fpname string) interface{} {
    57  	var ret interface{}
    58  	failpoint.Inject(fpname, func() {
    59  		ret = points.get(fpname)
    60  	})
    61  	return ret
    62  }
    63  
    64  // With enable failpoint and provide a value
    65  func With(fpname string, value interface{}) Finalizer {
    66  	if err := failpoint.Enable(failpath(fpname), "return(true)"); err != nil {
    67  		panic(err)
    68  	}
    69  	points.set(fpname, value)
    70  	return func() error { return Reset(fpname) }
    71  }
    72  
    73  // Reset disable failpoint and remove mock value
    74  func Reset(fpname string) error {
    75  	if err := failpoint.Disable(failpath(fpname)); err != nil {
    76  		return err
    77  	}
    78  	points.clr(fpname)
    79  	return nil
    80  }
    81  
    82  func failpath(fpname string) string {
    83  	type em struct{}
    84  	return path.Join(reflect.TypeOf(em{}).PkgPath(), fpname)
    85  }
    86