redefine the GetAgent method

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

test/e2e/e2e_agent_test.go

@@ -42,8 +42,9 @@ type PlannerAgent interface {
 
 type PlannerService interface {
 	RemoveSources() error
+	RemoveAgent(UUID string) error
 	GetSource() (*api.Source, error)
-	GetAgent() (*api.Agent, error)
+	GetAgent(credentialUrl string) (*api.Agent, error)
 }
 
 type plannerService struct {
@@ -260,7 +261,7 @@ func NewPlannerService(configPath string) (*plannerService, error) {
 	return &plannerService{c: c}, nil
 }
 
-func (s *plannerService) GetAgent() (*api.Agent, error) {
+func (s *plannerService) GetAgent(credentialUrl string) (*api.Agent, error) {
 	ctx := context.TODO()
 	res, err := s.c.ListAgentsWithResponse(ctx)
 	if err != nil || res.HTTPResponse.StatusCode != 200 {
@@ -271,9 +272,12 @@ func (s *plannerService) GetAgent() (*api.Agent, error) {
 		return nil, fmt.Errorf("No agents found")
 	}
 
-	nullUuid := uuid.UUID{}
+	//nullUuid := uuid.UUID{}
 	for _, agent := range *res.JSON200 {
-		if agent.Id != nullUuid.String() {
+		//if agent.Id != nullUuid.String() {
+		//	return &agent, nil
+		//}
+		if agent.CredentialUrl == credentialUrl {
 			return &agent, nil
 		}
 	}
@@ -302,6 +306,18 @@ func (s *plannerService) GetSource() (*api.Source, error) {
 	return nil, fmt.Errorf("No sources found")
 }
 
+func (s *plannerService) RemoveAgent(UUID string) error {
+	parsedUUID, err1 := uuid.Parse(UUID)
+	if err1 != nil {
+		return err1
+	}
+	_, err2 := s.c.DeleteAgentWithResponse(context.TODO(), parsedUUID)
+	if err2 != nil {
+		return err2
+	}
+	return nil
+}
+
 func (s *plannerService) RemoveSources() error {
 	_, err := s.c.DeleteSourcesWithResponse(context.TODO())
 	return err

test/e2e/e2e_test.go

@@ -40,7 +40,7 @@ var _ = Describe("e2e", func() {
 		}, "3m", "2s").Should(BeTrue())
 
 		Eventually(func() string {
-			s, err := svc.GetAgent(0)
+			s, err := svc.GetAgent(fmt.Sprintf("http://%s:3333", agentIP))
 			if err != nil {
 				return ""
 			}
@@ -57,7 +57,7 @@ var _ = Describe("e2e", func() {
 	})
 
 	AfterEach(func() {
-		s, err := svc.GetAgent(0)
+		s, err := svc.GetAgent(fmt.Sprintf("http://%s:3333", agentIP))
 		if err != nil {
 			s = nil
 		}
@@ -119,7 +119,7 @@ var _ = Describe("e2e", func() {
 			Expect(err).To(BeNil())
 			Expect(res.StatusCode).To(Equal(http.StatusNoContent))
 			Eventually(func() bool {
-				apiAgent, err := svc.GetAgent(0)
+				apiAgent, err := svc.GetAgent(fmt.Sprintf("http://%s:3333", agentIP))
 				if err != nil {
 					return false
 				}