119 lines
2.6 KiB
Go
119 lines
2.6 KiB
Go
package common
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/aws/aws-sdk-go/aws"
|
|
"github.com/aws/aws-sdk-go/service/ec2"
|
|
"github.com/mitchellh/multistep"
|
|
)
|
|
|
|
const (
|
|
privateIP = "10.0.0.1"
|
|
publicIP = "192.168.1.1"
|
|
publicDNS = "public.dns.test"
|
|
)
|
|
|
|
func TestSSHHost(t *testing.T) {
|
|
origSshHostSleepDuration := sshHostSleepDuration
|
|
defer func() { sshHostSleepDuration = origSshHostSleepDuration }()
|
|
sshHostSleepDuration = 0
|
|
|
|
var cases = []struct {
|
|
allowTries int
|
|
vpcId string
|
|
private bool
|
|
|
|
ok bool
|
|
wantHost string
|
|
}{
|
|
{1, "", false, true, publicDNS},
|
|
{1, "", true, true, privateIP},
|
|
{1, "vpc-id", false, true, publicIP},
|
|
{1, "vpc-id", true, true, privateIP},
|
|
{2, "", false, true, publicDNS},
|
|
{2, "", true, true, privateIP},
|
|
{2, "vpc-id", false, true, publicIP},
|
|
{2, "vpc-id", true, true, privateIP},
|
|
{3, "", false, false, ""},
|
|
{3, "", true, false, ""},
|
|
{3, "vpc-id", false, false, ""},
|
|
{3, "vpc-id", true, false, ""},
|
|
}
|
|
|
|
for _, c := range cases {
|
|
testSSHHost(t, c.allowTries, c.vpcId, c.private, c.ok, c.wantHost)
|
|
}
|
|
}
|
|
|
|
func testSSHHost(t *testing.T, allowTries int, vpcId string, private, ok bool, wantHost string) {
|
|
t.Logf("allowTries=%d vpcId=%s private=%t ok=%t wantHost=%q", allowTries, vpcId, private, ok, wantHost)
|
|
|
|
e := &fakeEC2Describer{
|
|
allowTries: allowTries,
|
|
vpcId: vpcId,
|
|
privateIP: privateIP,
|
|
publicIP: publicIP,
|
|
publicDNS: publicDNS,
|
|
}
|
|
|
|
f := SSHHost(e, private)
|
|
st := &multistep.BasicStateBag{}
|
|
st.Put("instance", &ec2.Instance{
|
|
InstanceId: aws.String("instance-id"),
|
|
})
|
|
|
|
host, err := f(st)
|
|
|
|
if e.tries > allowTries {
|
|
t.Fatalf("got %d ec2 DescribeInstances tries, want %d", e.tries, allowTries)
|
|
}
|
|
|
|
switch {
|
|
case ok && err != nil:
|
|
t.Fatalf("expected no error, got %+v", err)
|
|
case !ok && err == nil:
|
|
t.Fatalf("expected error, got none and host %s", host)
|
|
}
|
|
|
|
if host != wantHost {
|
|
t.Fatalf("got host %s, want %s", host, wantHost)
|
|
}
|
|
}
|
|
|
|
type fakeEC2Describer struct {
|
|
allowTries int
|
|
tries int
|
|
|
|
vpcId string
|
|
privateIP, publicIP, publicDNS string
|
|
}
|
|
|
|
func (d *fakeEC2Describer) DescribeInstances(in *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) {
|
|
d.tries++
|
|
|
|
instance := &ec2.Instance{
|
|
InstanceId: aws.String("instance-id"),
|
|
}
|
|
|
|
if d.vpcId != "" {
|
|
instance.VpcId = aws.String(d.vpcId)
|
|
}
|
|
|
|
if d.tries >= d.allowTries {
|
|
instance.PublicIpAddress = aws.String(d.publicIP)
|
|
instance.PrivateIpAddress = aws.String(d.privateIP)
|
|
instance.PublicDnsName = aws.String(d.publicDNS)
|
|
}
|
|
|
|
out := &ec2.DescribeInstancesOutput{
|
|
Reservations: []*ec2.Reservation{
|
|
{
|
|
Instances: []*ec2.Instance{instance},
|
|
},
|
|
},
|
|
}
|
|
|
|
return out, nil
|
|
}
|