go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/exec/execmock/intercept.go (about)

     1  // Copyright 2023 The LUCI 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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package execmock
    16  
    17  import (
    18  	"encoding/json"
    19  	"flag"
    20  	"fmt"
    21  	"log"
    22  	"os"
    23  	"reflect"
    24  	"strconv"
    25  
    26  	"go.chromium.org/luci/common/exec/internal/execmockctx"
    27  	"go.chromium.org/luci/common/exec/internal/execmockserver"
    28  )
    29  
    30  var server *execmockserver.Server
    31  
    32  func startServer() {
    33  	server = execmockserver.Start()
    34  }
    35  
    36  var execmockList = flag.Bool("execmock.list", false, "Lists the names of all registered runners and their ID.")
    37  
    38  const (
    39  	execmockRunnerIDEnv     = "EXECMOCK_RUNNER_ID"
    40  	execmockRunnerInputEnv  = "EXECMOCK_RUNNER_INPUT"
    41  	execmockRunnerOutputEnv = "EXECMOCK_RUNNER_OUTPUT"
    42  )
    43  
    44  // Intercept must be called from TestMain like:
    45  //
    46  //	func TestMain(m *testing.M) {
    47  //		execmock.Intercept()
    48  //		os.Exit(m.Run())
    49  //	}
    50  //
    51  // If process flags have not yet been parsed, this will call flag.Parse().
    52  func Intercept() {
    53  	runnerMu.Lock()
    54  	runnerRegistryMutable = false
    55  	registry := runnerRegistry
    56  	registryMeta := runnerRegistryMeta
    57  	runnerMu.Unlock()
    58  
    59  	if exitcode, intercepted := execmockserver.ClientIntercept(registry); intercepted {
    60  		os.Exit(exitcode)
    61  	}
    62  
    63  	if !flag.Parsed() {
    64  		flag.Parse()
    65  	}
    66  
    67  	if *execmockList {
    68  		fmt.Println("execmock Registered Runners:")
    69  		fmt.Println()
    70  		fmt.Println("<ID>: Type - Registration location")
    71  		for id, entry := range registry {
    72  			t := entry.Type()
    73  			meta := registryMeta[id]
    74  			if meta.file != "" {
    75  				fmt.Printf("  %d: Runner[%s, %s] - %s:%d\n", id, t.In(0), t.Out(0), meta.file, meta.line)
    76  			} else {
    77  				fmt.Printf("  %d: Runner[%s, %s]\n", id, t.In(0), t.Out(0))
    78  			}
    79  		}
    80  		fmt.Println()
    81  		fmt.Printf("To execute a single runner, set the envvar $%s to <ID>.\n", execmockRunnerIDEnv)
    82  		fmt.Printf("To provide input, write it in JSON to a file and set the file in $%s.\n", execmockRunnerInputEnv)
    83  		fmt.Printf("To see output, set the output file in $%s.\n", execmockRunnerOutputEnv)
    84  		fmt.Printf("After preparing the environment, run `go test`.")
    85  		os.Exit(0)
    86  	}
    87  
    88  	if idStr := os.Getenv(execmockRunnerIDEnv); idStr != "" {
    89  		inFilePath := os.Getenv(execmockRunnerInputEnv)
    90  		outFilePath := os.Getenv(execmockRunnerOutputEnv)
    91  
    92  		// prune all execmock envars from Env
    93  		os.Unsetenv(execmockRunnerIDEnv)
    94  		os.Unsetenv(execmockRunnerInputEnv)
    95  		os.Unsetenv(execmockRunnerOutputEnv)
    96  
    97  		id, err := strconv.ParseUint(idStr, 10, 64)
    98  		if err != nil {
    99  			log.Fatalf("execmock: $%s: not an id: %s", execmockRunnerIDEnv, err)
   100  		}
   101  
   102  		fn, ok := registry[id]
   103  		if !ok {
   104  			log.Fatalf("execmock: $%s: unknown Runner id: %d", execmockRunnerIDEnv, id)
   105  		}
   106  
   107  		inData := reflect.New(fn.Type().In(0)).Elem()
   108  		if inFilePath != "" {
   109  			inFile, err := os.Open(inFilePath)
   110  			if err != nil {
   111  				log.Fatalf("opening execmock.input: %s", err)
   112  			}
   113  			if err = json.NewDecoder(inFile).Decode(inData.Addr().Interface()); err != nil {
   114  				log.Fatalf("decoding execmock.input: %s", err)
   115  			}
   116  		}
   117  
   118  		results := fn.Call([]reflect.Value{inData})
   119  
   120  		if outFilePath != "" {
   121  			outFile, err := os.Create(outFilePath)
   122  			if err != nil {
   123  				log.Fatalf("opening execmock.output: %s", err)
   124  			}
   125  			toEnc := struct {
   126  				Error string
   127  				Data  any
   128  			}{Data: results[0].Interface()}
   129  			if errVal := results[2]; !errVal.IsNil() {
   130  				toEnc.Error = errVal.Interface().(error).Error()
   131  			}
   132  			if err = json.NewEncoder(outFile).Encode(toEnc); err != nil {
   133  				log.Fatalf("encoding execmock.output: %s", err)
   134  			}
   135  		}
   136  
   137  		os.Exit(int(results[1].Int()))
   138  	}
   139  
   140  	// This is the real `go test` invocation.
   141  	execmockctx.EnableMockingForThisProcess(getMocker)
   142  	startServer()
   143  }