From 3dc72f3f353f7c257744ced0abf9b78381aef191 Mon Sep 17 00:00:00 2001
From: Paul Kuruvilla <rohitpaulk@gmail.com>
Date: Mon, 4 Mar 2024 17:30:21 +0000
Subject: [PATCH] Refactor test command and codecrafters client

---
 internal/commands/test.go             | 45 +++++++++++----------------
 internal/utils/codecrafters_client.go | 37 +++++++++++++++++++---
 2 files changed, 50 insertions(+), 32 deletions(-)

diff --git a/internal/commands/test.go b/internal/commands/test.go
index 71b98e4..8820fbd 100644
--- a/internal/commands/test.go
+++ b/internal/commands/test.go
@@ -17,7 +17,6 @@ import (
 	logstream_consumer "github.com/codecrafters-io/logstream/consumer"
 	"github.com/fatih/color"
 	"github.com/getsentry/sentry-go"
-	wordwrap "github.com/mitchellh/go-wordwrap"
 	cp "github.com/otiai10/copy"
 	"github.com/rs/zerolog"
 )
@@ -89,7 +88,7 @@ func TestCommand(ctx context.Context) (err error) {
 	}
 
 	// Place this before the push so that it "feels" fast
-	fmt.Println("Running tests on your codebase. Streaming logs...")
+	fmt.Println("Initiating test run...")
 
 	logger.Debug().Msg("push changes")
 
@@ -111,29 +110,15 @@ func TestCommand(ctx context.Context) (err error) {
 
 	logger.Debug().Interface("response", createSubmissionResponse).Msg("submission created")
 
-	if createSubmissionResponse.OnInitSuccessMessage != "" {
+	for _, message := range createSubmissionResponse.OnInitMessages {
 		fmt.Println("")
-
-		wrapped := wordwrap.WrapString(createSubmissionResponse.OnInitSuccessMessage, 79)
-		for _, line := range strings.Split(wrapped, "\n") {
-			fmt.Printf("\033[1;92m%s\033[0m\n", line)
-		}
-	}
-
-	if createSubmissionResponse.OnInitWarningMessage != "" {
-		fmt.Println("")
-
-		wrapped := wordwrap.WrapString(createSubmissionResponse.OnInitWarningMessage, 79)
-		for _, line := range strings.Split(wrapped, "\n") {
-			fmt.Printf("\033[31m%s\033[0m\n", line)
-		}
+		message.Print()
 	}
 
-	fmt.Println("")
-
 	if createSubmissionResponse.BuildLogstreamURL != "" {
 		logger.Debug().Msg("stream build logs")
 
+		fmt.Println("")
 		err = streamLogs(createSubmissionResponse.BuildLogstreamURL)
 		if err != nil {
 			return fmt.Errorf("stream build logs: %w", err)
@@ -165,10 +150,6 @@ func TestCommand(ctx context.Context) (err error) {
 			os.Exit(0)
 		case "success":
 			time.Sleep(1 * time.Second) // The delay in-between build and test logs is usually 5-10 seconds, so let's buy some time
-
-			fmt.Println("")
-			fmt.Println("Running tests. Logs should appear shortly...")
-			fmt.Println("")
 		default:
 			red := color.New(color.FgRed).SprintFunc()
 
@@ -180,6 +161,10 @@ func TestCommand(ctx context.Context) (err error) {
 
 	logger.Debug().Msg("stream logs")
 
+	fmt.Println("")
+	fmt.Println("Running tests. Logs should appear shortly...")
+	fmt.Println("")
+
 	err = streamLogs(createSubmissionResponse.LogstreamURL)
 	if err != nil {
 		return fmt.Errorf("stream logs: %w", err)
@@ -200,13 +185,19 @@ func TestCommand(ctx context.Context) (err error) {
 
 	logger.Debug().Interface("response", createSubmissionResponse).Msg("submission fetched")
 
-	fmt.Println("")
-
 	switch fetchSubmissionResponse.Status {
 	case "failure":
-		fmt.Println(createSubmissionResponse.OnFailureMessage)
+		for _, message := range createSubmissionResponse.OnFailureMessages {
+			fmt.Println("")
+			message.Print()
+		}
 	case "success":
-		fmt.Println(createSubmissionResponse.OnSuccessMessage)
+		for _, message := range createSubmissionResponse.OnSuccessMessages {
+			fmt.Println("")
+			message.Print()
+		}
+	default:
+		fmt.Println("")
 	}
 
 	if fetchSubmissionResponse.IsError {
diff --git a/internal/utils/codecrafters_client.go b/internal/utils/codecrafters_client.go
index c475e4f..0f40a71 100644
--- a/internal/utils/codecrafters_client.go
+++ b/internal/utils/codecrafters_client.go
@@ -4,12 +4,40 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"strings"
 	"time"
 
 	retry "github.com/avast/retry-go"
 	"github.com/levigross/grequests"
+	"github.com/mitchellh/go-wordwrap"
 )
 
+type Message struct {
+	Text  string `json:"text"`
+	Color string `json:"color"`
+}
+
+func (m Message) Print() {
+	wrapped := wordwrap.WrapString(m.Text, 79)
+
+	lineFormat := "%s\n"
+
+	switch m.Color {
+	case "red":
+		lineFormat = "\033[31m%s\033[0m\n"
+	case "green":
+		lineFormat = "\033[32m%s\033[0m\n"
+	case "yellow":
+		lineFormat = "\033[33m%s\033[0m\n"
+	case "blue":
+		lineFormat = "\033[34m%s\033[0m\n"
+	}
+
+	for _, line := range strings.Split(wrapped, "\n") {
+		fmt.Printf(lineFormat, line)
+	}
+}
+
 type CreateSubmissionResponse struct {
 	Id string `json:"id"`
 
@@ -22,11 +50,10 @@ type CreateSubmissionResponse struct {
 	// LogstreamURL contains test logs.
 	LogstreamURL string `json:"logstream_url"`
 
-	// Messages to be displayed to the user
-	OnSuccessMessage     string `json:"on_success_message"`
-	OnFailureMessage     string `json:"on_failure_message"`
-	OnInitSuccessMessage string `json:"on_init_success_message"`
-	OnInitWarningMessage string `json:"on_init_warning_message"`
+	// Messages to be displayed to the user at various stages of the submission lifecycle
+	OnInitMessages    []Message `json:"on_init_messages"`
+	OnSuccessMessages []Message `json:"on_success_messages"`
+	OnFailureMessages []Message `json:"on_failure_messages"`
 
 	// IsError is true when the submission failed to be created, and ErrorMessage is the human-friendly error message
 	IsError      bool   `json:"is_error"`