common_test.go

Documentation: github.com/GoogleCloudPlatform/cloudsql-proxy/tests

     1  // Copyright 2015 Google Inc. All Rights Reserved.
     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 tests contains end to end tests meant to verify the Cloud SQL Auth proxy
    16  // works as expected when executed as a binary.
    17  //
    18  // Required flags:
    19  //    -mysql_conn_name, -db_user, -db_pass
    20  package tests
    21  
    22  import (
    23  	"bufio"
    24  	"bytes"
    25  	"context"
    26  	"flag"
    27  	"fmt"
    28  	"io"
    29  	"io/ioutil"
    30  	"log"
    31  	"os"
    32  	"os/exec"
    33  	"path"
    34  	"runtime"
    35  	"strings"
    36  	"testing"
    37  )
    38  
    39  var (
    40  	binPath = ""
    41  )
    42  
    43  func TestMain(m *testing.M) {
    44  	flag.Parse()
    45  	// compile the proxy as a binary
    46  	var err error
    47  	binPath, err = compileProxy()
    48  	if err != nil {
    49  		log.Fatalf("failed to compile proxy: %s", err)
    50  	}
    51  	// Run tests and cleanup
    52  	rtn := m.Run()
    53  	os.RemoveAll(binPath)
    54  
    55  	os.Exit(rtn)
    56  }
    57  
    58  // compileProxy compiles the binary into a temporary directory, and returns the path to the file or any error that occured.
    59  func compileProxy() (string, error) {
    60  	// get path of the cmd pkg
    61  	_, f, _, ok := runtime.Caller(0)
    62  	if !ok {
    63  		return "", fmt.Errorf("failed to find cmd pkg")
    64  	}
    65  	projRoot := path.Dir(path.Dir(f)) // cd ../..
    66  	pkgPath := path.Join(projRoot, "cmd", "cloud_sql_proxy")
    67  	// compile the proxy into a tmp directory
    68  	tmp, err := ioutil.TempDir("", "")
    69  	if err != nil {
    70  		return "", fmt.Errorf("failed to create temp dir: %s", err)
    71  	}
    72  
    73  	b := path.Join(tmp, "cloud_sql_proxy")
    74  
    75  	if runtime.GOOS == "windows" {
    76  		b += ".exe"
    77  	}
    78  
    79  	cmd := exec.Command("go", "build", "-o", b, pkgPath)
    80  	out, err := cmd.CombinedOutput()
    81  	if err != nil {
    82  		return "", fmt.Errorf("failed to run 'go build': %w \n %s", err, out)
    83  	}
    84  	return b, nil
    85  }
    86  
    87  // proxyExec represents an execution of the Cloud SQL proxy.
    88  type ProxyExec struct {
    89  	Out io.ReadCloser
    90  
    91  	cmd     *exec.Cmd
    92  	cancel  context.CancelFunc
    93  	closers []io.Closer
    94  	done    chan bool // closed once the cmd is completed
    95  	err     error
    96  }
    97  
    98  // StartProxy returns a proxyExec representing a running instance of the proxy.
    99  func StartProxy(ctx context.Context, args ...string) (*ProxyExec, error) {
   100  	var err error
   101  	ctx, cancel := context.WithCancel(ctx)
   102  	p := ProxyExec{
   103  		cmd:    exec.CommandContext(ctx, binPath, args...),
   104  		cancel: cancel,
   105  		done:   make(chan bool),
   106  	}
   107  	pr, pw, err := os.Pipe()
   108  	if err != nil {
   109  		return nil, fmt.Errorf("unable to open stdout pipe: %w", err)
   110  	}
   111  	defer pw.Close()
   112  	p.Out, p.cmd.Stdout, p.cmd.Stderr = pr, pw, pw
   113  	p.closers = append(p.closers, pr)
   114  	if err := p.cmd.Start(); err != nil {
   115  		defer p.Close()
   116  		return nil, fmt.Errorf("unable to start cmd: %w", err)
   117  	}
   118  	// when process is complete, mark as finished
   119  	go func() {
   120  		defer close(p.done)
   121  		p.err = p.cmd.Wait()
   122  	}()
   123  	return &p, nil
   124  }
   125  
   126  // Stop sends the pskill signal to the proxy and returns.
   127  func (p *ProxyExec) Kill() {
   128  	p.cancel()
   129  }
   130  
   131  // Waits until the execution is completed and returns any error.
   132  func (p *ProxyExec) Wait() error {
   133  	select {
   134  	case <-p.done:
   135  		return p.err
   136  	}
   137  }
   138  
   139  // Stop sends the pskill signal to the proxy and returns.
   140  func (p *ProxyExec) Done() bool {
   141  	select {
   142  	case <-p.done:
   143  		return true
   144  	default:
   145  	}
   146  	return false
   147  }
   148  
   149  // Close releases any resources assotiated with the instance.
   150  func (p *ProxyExec) Close() {
   151  	p.cancel()
   152  	for _, c := range p.closers {
   153  		c.Close()
   154  	}
   155  }
   156  
   157  // WaitForServe waits until the proxy ready to serve traffic. Returns any output from the proxy
   158  // while starting or any errors experienced before the proxy was ready to server.
   159  func (p *ProxyExec) WaitForServe(ctx context.Context) (output string, err error) {
   160  	// Watch for the "Ready for new connections" to indicate the proxy is listening
   161  	buf, in, errCh := new(bytes.Buffer), bufio.NewReader(p.Out), make(chan error, 1)
   162  	go func() {
   163  		defer close(errCh)
   164  		for {
   165  			// if ctx is finished, stop processing
   166  			select {
   167  			case <-ctx.Done():
   168  				return
   169  			default:
   170  			}
   171  			s, err := in.ReadString('\n')
   172  			if err != nil {
   173  				errCh <- err
   174  				return
   175  			}
   176  			buf.WriteString(s)
   177  			if strings.Contains(s, "Ready for new connections") {
   178  				errCh <- nil
   179  				return
   180  			}
   181  		}
   182  	}()
   183  	// Wait for either the background thread of the context to complete
   184  	select {
   185  	case <-ctx.Done():
   186  		return buf.String(), fmt.Errorf("context done: %w", ctx.Err())
   187  	case err := <-errCh:
   188  		if err != nil {
   189  			return buf.String(), fmt.Errorf("proxy start failed: %w", err)
   190  		}
   191  	}
   192  	return buf.String(), nil
   193  }
   194
View as plain text