fixes

internal/app/rest_server.go

@@ -8,7 +8,9 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/google/uuid"
@@ -29,6 +31,7 @@ var (
 )
 
 type RestServer struct {
+	ok        atomic.Bool
 	stuber    *stuber.Budgerigar
 	convertor *yaml2json.Convertor
 	caser     cases.Caser
@@ -113,6 +116,12 @@ func (h *RestServer) Liveness(w http.ResponseWriter, _ *http.Request) {
 }
 
 func (h *RestServer) Readiness(w http.ResponseWriter, _ *http.Request) {
+	if !h.ok.Load() {
+		w.WriteHeader(http.StatusServiceUnavailable)
+
+		return
+	}
+
 	w.Header().Set("Content-Type", "application/json")
 
 	if err := json.NewEncoder(w).Encode(rest.MessageOK{Message: "ok", Time: time.Now()}); err != nil {
@@ -276,49 +285,69 @@ func (h *RestServer) writeResponseError(err error, w http.ResponseWriter) {
 	}
 }
 
+// readStubs reads all the stubs from the given directory and its subdirectories,
+// and adds them to the server's stub store.
+// The stub files can be in yaml or json format.
+// If a file is in yaml format, it will be converted to json format.
 func (h *RestServer) readStubs(path string) {
+	defer h.ok.Store(true)
+
 	files, err := os.ReadDir(path)
 	if err != nil {
-		log.Printf("Can't read stub from %s. %v\n", path, err)
-
+		log.Printf("can't read stubs from %s: %v", path, err)
 		return
 	}
 
 	for _, file := range files {
+		// If the file is a directory, recursively read its stubs.
 		if file.IsDir() {
-			h.readStubs(path + "/" + file.Name())
-
+			h.readStubs(filepath.Join(path, file.Name()))
 			continue
 		}
 
-		byt, err := os.ReadFile(path + "/" + file.Name())
-		if err != nil {
-			log.Printf("Error when reading file %s. %v. skipping...", file.Name(), err)
-
+		// Only process files with yaml or yml extensions.
+		if !strings.HasSuffix(file.Name(), ".yaml") && !strings.HasSuffix(file.Name(), ".yml") {
 			continue
 		}
 
-		if strings.HasSuffix(file.Name(), ".yaml") || strings.HasSuffix(file.Name(), ".yml") {
-			byt, err = h.convertor.Execute(file.Name(), byt)
-			if err != nil {
-				log.Printf("Error when unmarshalling file %s. %v. skipping...", file.Name(), err)
-
-				continue
-			}
+		// Read the stub file and add it to the server's stub store.
+		stubs, err := h.readStub(filepath.Join(path, file.Name()))
+		if err != nil {
+			log.Printf("cant read stubs from %s: %v", file.Name(), err)
+			continue
 		}
+		h.stuber.PutMany(stubs...)
+	}
+}
 
-		var storageStubs []*stuber.Stub
-
-		if err = jsondecoder.UnmarshalSlice(byt, &storageStubs); err != nil {
-			log.Printf("Error when unmarshalling file %s. %v %v. skipping...", file.Name(), string(byt), err)
+// readStub reads a stub file and returns a slice of stubs.
+// The stub file can be in yaml or json format.
+// If the file is in yaml format, it will be converted to json format.
+func (h *RestServer) readStub(path string) ([]*stuber.Stub, error) {
+	// Read the file
+	byt, err := os.ReadFile(path)
+	if err != nil {
+		return nil, fmt.Errorf("error when reading file %s: %v", path, err)
+	}
 
-			continue
+	// If the file is in yaml format, convert it to json format
+	if strings.HasSuffix(path, ".yaml") || strings.HasSuffix(path, ".yml") {
+		byt, err = h.convertor.Execute(path, byt)
+		if err != nil {
+			return nil, fmt.Errorf("error when unmarshalling file %s: %v", path, err)
 		}
+	}
 
-		h.stuber.PutMany(storageStubs...)
+	// Unmarshal the json into a slice of stubs
+	var stubs []*stuber.Stub
+	if err := jsondecoder.UnmarshalSlice(byt, &stubs); err != nil {
+		return nil, fmt.Errorf("error when unmarshalling file %s: %v %s", path, string(byt), err)
 	}
+
+	return stubs, nil
 }
 
+// validateStub validates if the stub is valid or not.
 func validateStub(stub *stuber.Stub) error {
 	if stub.Service == "" {
 		return ErrServiceIsMissing
@@ -341,12 +370,8 @@ func validateStub(stub *stuber.Stub) error {
 		return fmt.Errorf("input cannot be empty")
 	}
 
-	// TODO: validate all input case
-
 	if stub.Output.Error == "" && stub.Output.Data == nil && stub.Output.Code == nil {
-		// fixme
-		//nolint:goerr113,perfsprint
-		return fmt.Errorf("output can't be empty")
+		return fmt.Errorf("output cannot be empty")
 	}
 
 	return nil

main.go

@@ -89,19 +89,12 @@ func main() {
 		}
 	}
 
-	chReady := make(chan struct{})
-	defer close(chReady)
-
 	// Run the admin stub server in a separate goroutine.
 	//
 	// This goroutine runs the REST server that serves the stub files.
 	// It waits for the ready signal from the gRPC server goroutine.
 	// Once the gRPC server is ready, it starts the admin stub server.
 	go func() {
-		<-chReady
-
-		zerolog.Ctx(ctx).Info().Msg("gRPC server is ready to accept requests")
-
 		stub.RunRestServer(ctx, *stubPath, builder.Config(), builder.Reflector())
 	}()
 
@@ -142,7 +135,7 @@ func main() {
 
 		// If the server is in the "SERVING" state, send a signal to the chReady channel.
 		if check.GetStatus() == healthv1.HealthCheckResponse_SERVING {
-			chReady <- struct{}{}
+			zerolog.Ctx(ctx).Info().Msg("gRPC server is ready to accept requests")
 		}
 	}()
 

protoc-gen-gripmock/server.tmpl

@@ -3,10 +3,9 @@ package main
 
 import (
 	"context"
-	"errors"
+	"time"
 	"slices"
 	"fmt"
-	"io"
 	"log"
 	"net"
 	"net/http"
@@ -63,19 +62,52 @@ func main() {
             grpccontext.StreamInterceptor(builder.Logger()),
         }...),
     )
+
+	healthcheck := health.NewServer()
+	healthcheck.SetServingStatus("", healthgrpc.HealthCheckResponse_NOT_SERVING)
+
 	{{ range .Services }}
 	{{ template "register_services" . }}
 	{{ end }}
-
-	healthgrpc.RegisterHealthServer(s, health.NewServer())
+	healthgrpc.RegisterHealthServer(s, healthcheck)
     reflection.Register(s)
 
 	builder.Logger().Info().
-	    Str("addr", fmt.Sprintf("%s://%s", builder.Config().GRPCNetwork, builder.Config().GRPCAddr)).
+		Str("addr", builder.Config().GRPCAddr).
+	    Str("network", builder.Config().GRPCNetwork).
 	    Msg("Serving gRPC")
 
+	go func () {
+		api, err := sdk.NewClientWithResponses(
+			fmt.Sprintf("http://%s/api", builder.Config().HTTPAddr),
+        	sdk.WithHTTPClient(http.DefaultClient))
+		if err != nil {
+			return
+		}	
+
+		ctx, cancel := context.WithTimeout(ctx, 120 * time.Second)
+		defer cancel()
+
+		tick := time.NewTicker(250 * time.Millisecond)
+		defer tick.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-tick.C:
+				resp, err := api.ReadinessWithResponse(ctx)
+				if err == nil && resp.JSON200 != nil {
+					healthcheck.SetServingStatus("", healthgrpc.HealthCheckResponse_SERVING)
+
+					return
+				}
+			}
+		}
+	}()
+
 	if err := s.Serve(lis); err != nil {
-		builder.Logger().Fatal().Err(err).Msg("server ended")
+		builder.Logger().Fatal().Err(err).Msg("failed to serve")
 	}
 }