From ff84803546c958c3f981d4396e6727d58795a994 Mon Sep 17 00:00:00 2001
From: Guillaume Tardif <guillaume.tardif@docker.com>
Date: Thu, 19 Nov 2020 13:06:39 +0100
Subject: [PATCH] compose ps implementation

Signed-off-by: Guillaume Tardif <guillaume.tardif@docker.com>
---
 local/compose.go      | 84 ++++++++++++++++++++++++++++---------------
 local/compose_test.go | 41 +++++++++++++++++++++
 2 files changed, 96 insertions(+), 29 deletions(-)

diff --git a/local/compose.go b/local/compose.go
index de6fe593d..f32f16e79 100644
--- a/local/compose.go
+++ b/local/compose.go
@@ -247,19 +247,51 @@ func (s *local) Ps(ctx context.Context, projectName string) ([]compose.ServiceSt
 	if err != nil {
 		return nil, err
 	}
-	var status []compose.ServiceStatus
-	for _, c := range list {
-		// TODO group by service
-		status = append(status, compose.ServiceStatus{
-			ID:         c.ID,
-			Name:       c.Labels[serviceLabel],
-			Replicas:   0,
-			Desired:    0,
-			Ports:      nil,
-			Publishers: nil,
+	return containersToServiceStatus(list)
+}
+
+func containersToServiceStatus(containers []moby.Container) ([]compose.ServiceStatus, error) {
+	containersByLabel, keys, err := groupContainerByLabel(containers, serviceLabel)
+	if err != nil {
+		return nil, err
+	}
+	var services []compose.ServiceStatus
+	for _, service := range keys {
+		containers := containersByLabel[service]
+		runnningContainers := []moby.Container{}
+		for _, container := range containers {
+			if container.State == "running" {
+				runnningContainers = append(runnningContainers, container)
+			}
+		}
+		services = append(services, compose.ServiceStatus{
+			ID:       service,
+			Name:     service,
+			Desired:  len(containers),
+			Replicas: len(runnningContainers),
 		})
 	}
-	return status, nil
+	return services, nil
+}
+
+func groupContainerByLabel(containers []moby.Container, labelName string) (map[string][]moby.Container, []string, error) {
+	containersByLabel := map[string][]moby.Container{}
+	keys := []string{}
+	for _, c := range containers {
+		label, ok := c.Labels[labelName]
+		if !ok {
+			return nil, nil, fmt.Errorf("No label %q set on container %q of compose project", labelName, c.ID)
+		}
+		labelContainers, ok := containersByLabel[label]
+		if !ok {
+			labelContainers = []moby.Container{}
+			keys = append(keys, label)
+		}
+		labelContainers = append(labelContainers, c)
+		containersByLabel[label] = labelContainers
+	}
+	sort.Strings(keys)
+	return containersByLabel, keys, nil
 }
 
 func (s *local) List(ctx context.Context, projectName string) ([]compose.Stack, error) {
@@ -274,35 +306,29 @@ func (s *local) List(ctx context.Context, projectName string) ([]compose.Stack,
 }
 
 func containersToStacks(containers []moby.Container) ([]compose.Stack, error) {
-	statusesByProject := map[string][]string{}
-	keys := []string{}
-	for _, c := range containers {
-		project, ok := c.Labels[projectLabel]
-		if !ok {
-			return nil, fmt.Errorf("No label %q set on container %q of compose project", serviceLabel, c.ID)
-		}
-		projectStatuses, ok := statusesByProject[project]
-		if !ok {
-			projectStatuses = []string{}
-			keys = append(keys, project)
-		}
-		projectStatuses = append(projectStatuses, c.State)
-		statusesByProject[project] = projectStatuses
+	containersByLabel, keys, err := groupContainerByLabel(containers, projectLabel)
+	if err != nil {
+		return nil, err
 	}
-
-	sort.Strings(keys)
 	var projects []compose.Stack
 	for _, project := range keys {
-		statuses := statusesByProject[project]
 		projects = append(projects, compose.Stack{
 			ID:     project,
 			Name:   project,
-			Status: combinedStatus(statuses),
+			Status: combinedStatus(containerToState(containersByLabel[project])),
 		})
 	}
 	return projects, nil
 }
 
+func containerToState(containers []moby.Container) []string {
+	statuses := []string{}
+	for _, c := range containers {
+		statuses = append(statuses, c.State)
+	}
+	return statuses
+}
+
 func combinedStatus(statuses []string) string {
 	nbByStatus := map[string]int{}
 	keys := []string{}
diff --git a/local/compose_test.go b/local/compose_test.go
index 1fca806d9..74de0745b 100644
--- a/local/compose_test.go
+++ b/local/compose_test.go
@@ -61,6 +61,47 @@ func TestContainersToStacks(t *testing.T) {
 	})
 }
 
+func TestContainersToServiceStatus(t *testing.T) {
+	containers := []types.Container{
+		{
+			ID:     "c1",
+			State:  "running",
+			Labels: map[string]string{serviceLabel: "service1"},
+		},
+		{
+			ID:     "c2",
+			State:  "exited",
+			Labels: map[string]string{serviceLabel: "service1"},
+		},
+		{
+			ID:     "c3",
+			State:  "running",
+			Labels: map[string]string{serviceLabel: "service1"},
+		},
+		{
+			ID:     "c4",
+			State:  "running",
+			Labels: map[string]string{serviceLabel: "service2"},
+		},
+	}
+	services, err := containersToServiceStatus(containers)
+	assert.NilError(t, err)
+	assert.DeepEqual(t, services, []compose.ServiceStatus{
+		{
+			ID:       "service1",
+			Name:     "service1",
+			Replicas: 2,
+			Desired:  3,
+		},
+		{
+			ID:       "service2",
+			Name:     "service2",
+			Replicas: 1,
+			Desired:  1,
+		},
+	})
+}
+
 func TestStacksMixedStatus(t *testing.T) {
 	assert.Equal(t, combinedStatus([]string{"running"}), "running(1)")
 	assert.Equal(t, combinedStatus([]string{"running", "running", "running"}), "running(3)")