github.com/opentofu/opentofu@v1.7.1/internal/communicator/communicator_mock.go (about)

     1  // Copyright (c) The OpenTofu Authors
     2  // SPDX-License-Identifier: MPL-2.0
     3  // Copyright (c) 2023 HashiCorp, Inc.
     4  // SPDX-License-Identifier: MPL-2.0
     5  
     6  package communicator
     7  
     8  import (
     9  	"bytes"
    10  	"fmt"
    11  	"io"
    12  	"strings"
    13  	"time"
    14  
    15  	"github.com/opentofu/opentofu/internal/communicator/remote"
    16  	"github.com/opentofu/opentofu/internal/provisioners"
    17  )
    18  
    19  // MockCommunicator is an implementation of Communicator that can be used for tests.
    20  type MockCommunicator struct {
    21  	RemoteScriptPath string
    22  	Commands         map[string]bool
    23  	Uploads          map[string]string
    24  	UploadScripts    map[string]string
    25  	UploadDirs       map[string]string
    26  	CommandFunc      func(*remote.Cmd) error
    27  	DisconnectFunc   func() error
    28  	ConnTimeout      time.Duration
    29  }
    30  
    31  // Connect implementation of communicator.Communicator interface
    32  func (c *MockCommunicator) Connect(o provisioners.UIOutput) error {
    33  	return nil
    34  }
    35  
    36  // Disconnect implementation of communicator.Communicator interface
    37  func (c *MockCommunicator) Disconnect() error {
    38  	if c.DisconnectFunc != nil {
    39  		return c.DisconnectFunc()
    40  	}
    41  	return nil
    42  }
    43  
    44  // Timeout implementation of communicator.Communicator interface
    45  func (c *MockCommunicator) Timeout() time.Duration {
    46  	if c.ConnTimeout != 0 {
    47  		return c.ConnTimeout
    48  	}
    49  	return time.Duration(5 * time.Second)
    50  }
    51  
    52  // ScriptPath implementation of communicator.Communicator interface
    53  func (c *MockCommunicator) ScriptPath() string {
    54  	return c.RemoteScriptPath
    55  }
    56  
    57  // Start implementation of communicator.Communicator interface
    58  func (c *MockCommunicator) Start(r *remote.Cmd) error {
    59  	r.Init()
    60  
    61  	if c.CommandFunc != nil {
    62  		return c.CommandFunc(r)
    63  	}
    64  
    65  	if !c.Commands[r.Command] {
    66  		return fmt.Errorf("Command not found!")
    67  	}
    68  
    69  	r.SetExitStatus(0, nil)
    70  
    71  	return nil
    72  }
    73  
    74  // Upload implementation of communicator.Communicator interface
    75  func (c *MockCommunicator) Upload(path string, input io.Reader) error {
    76  	f, ok := c.Uploads[path]
    77  	if !ok {
    78  		return fmt.Errorf("Path %q not found!", path)
    79  	}
    80  
    81  	var buf bytes.Buffer
    82  	buf.ReadFrom(input)
    83  	content := strings.TrimSpace(buf.String())
    84  
    85  	f = strings.TrimSpace(f)
    86  	if f != content {
    87  		return fmt.Errorf("expected: %q\n\ngot: %q\n", f, content)
    88  	}
    89  
    90  	return nil
    91  }
    92  
    93  // UploadScript implementation of communicator.Communicator interface
    94  func (c *MockCommunicator) UploadScript(path string, input io.Reader) error {
    95  	c.Uploads = c.UploadScripts
    96  	return c.Upload(path, input)
    97  }
    98  
    99  // UploadDir implementation of communicator.Communicator interface
   100  func (c *MockCommunicator) UploadDir(dst string, src string) error {
   101  	v, ok := c.UploadDirs[src]
   102  	if !ok {
   103  		return fmt.Errorf("Directory not found!")
   104  	}
   105  
   106  	if v != dst {
   107  		return fmt.Errorf("expected: %q\n\ngot: %q\n", v, dst)
   108  	}
   109  
   110  	return nil
   111  }