Add test case for vm reboot

Signed-off-by: Aviel Segev <[email protected]>

test/e2e/e2e_agent_test.go

@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/kubev2v/migration-planner/api/v1alpha1"
@@ -39,6 +40,7 @@ type PlannerAgent interface {
 	Run() error
 	Login(url string, user string, pass string) (*http.Response, error)
 	Version() (string, error)
+	Restart() error
 	Remove() error
 	GetIp() (string, error)
 	IsServiceRunning(string, string) bool
@@ -221,6 +223,61 @@ func (p *plannerAgentLibvirt) RestartService() error {
 	return nil
 }
 
+func (p *plannerAgentLibvirt) Restart() error {
+	domain, err := p.con.LookupDomainByName(p.name)
+	if err != nil {
+		return fmt.Errorf("failed to find vm: %w", err)
+	}
+
+	defer func() {
+		_ = domain.Free()
+	}()
+
+	// power off the vm
+	if err = domain.Shutdown(); err != nil {
+		return fmt.Errorf("failed to shutdown vm: %w", err)
+	}
+
+	// Wait for shutdown with timeout
+	if err = waitForDomainState(30*time.Second, domain, libvirt.DOMAIN_SHUTOFF); err != nil {
+		return fmt.Errorf("failed to reach shutdown state: %w", err)
+	}
+
+	// start the vm
+	err = domain.Create()
+	if err != nil {
+		return fmt.Errorf("failed to start vm: %w", err)
+	}
+
+	// Wait for startup with timeout
+	if err = waitForDomainState(30*time.Second, domain, libvirt.DOMAIN_RUNNING); err != nil {
+		return fmt.Errorf("failed to reach running state: %w", err)
+	}
+
+	return nil
+}
+
+func waitForDomainState(duration time.Duration, domain *libvirt.Domain, desiredState libvirt.DomainState) error {
+	timeout := time.After(duration)
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-timeout:
+			return fmt.Errorf("timed out waiting for desired state")
+		case <-ticker.C:
+			state, _, err := domain.GetState()
+			if err != nil {
+				return fmt.Errorf("failed to get VM state: %w", err)
+			}
+			if state == desiredState {
+				return nil
+			}
+		}
+	}
+}
+
 func (p *plannerAgentLibvirt) Remove() error {
 	defer p.con.Close()
 	domain, err := p.con.LookupDomainByName(p.name)

test/e2e/e2e_test.go

@@ -123,4 +123,33 @@ var _ = Describe("e2e", func() {
 			}, "1m", "2s").Should(BeTrue())
 		})
 	})
+
+	Context("Edge cases", func() {
+		It("VM reboot", func() {
+			res, err := agent.Login(fmt.Sprintf("https://%s:8989/sdk", systemIP), "core", "123456")
+			Expect(err).To(BeNil())
+			Expect(res.StatusCode).To(Equal(http.StatusNoContent))
+
+			// Restarting the VM
+			err = agent.Restart()
+			Expect(err).To(BeNil())
+
+			// Remove old ssh key
+			err = RemoveSSHKey(agentIP)
+			Expect(err).To(BeNil())
+
+			// Check that planner-agent service is running
+			Eventually(func() bool {
+				return agent.IsServiceRunning(agentIP, "planner-agent")
+			}, "3m").Should(BeTrue())
+
+			Eventually(func() bool {
+				source, err := svc.GetSource(source.Id)
+				if err != nil {
+					return false
+				}
+				return source.Agent.Status == v1alpha1.AgentStatusUpToDate
+			}, "2m").Should(BeTrue())
+		})
+	})
 })

test/e2e/e2e_utils_test.go

@@ -117,6 +117,25 @@ func CreateVm(c *libvirt.Connect) error {
 	return nil
 }
 
+func RemoveSSHKey(ip string) error {
+	// Check if the key exists in known_hosts
+	cmd := exec.Command("ssh-keygen", "-F", ip)
+	var Stdout bytes.Buffer
+	cmd.Stdout = &Stdout
+
+	if err := cmd.Run(); err != nil || Stdout.Len() == 0 {
+		return nil
+	}
+
+	// Key found, remove it
+	removeCmd := exec.Command("ssh-keygen", "-R", ip)
+	if err := removeCmd.Run(); err != nil {
+		return fmt.Errorf("failed to remove SSH key for %s: %v", ip, err)
+	}
+
+	return nil
+}
+
 func RunCommand(ip string, command string) (string, error) {
 	sshCmd := exec.Command("sshpass", "-p", "123456", "ssh", "-o", "StrictHostKeyChecking=no", fmt.Sprintf("core@%s", ip), command)