Add new methods to zrouter (#110)

* Add new methods to zrouter

* fix Mount

* minor change

* minor fix

* minor fix

* fix flaky tests

Author: lucaslopezf

pkg/zcache/local_cache_test.go

@@ -93,7 +93,8 @@ func (suite *LocalCacheTestSuite) TestCleanupProcess() {
 		CleanupProcess: CleanupProcess{
 			Interval: cleanupInterval,
 		},
-		MetricServer: metrics.NewTaskMetrics("", "", "appname")})
+		MetricServer: metrics.NewTaskMetrics("", "", "appname"),
+	})
 	suite.NoError(err)
 
 	ctx := context.Background()
@@ -103,12 +104,29 @@ func (suite *LocalCacheTestSuite) TestCleanupProcess() {
 	err = cache.Set(ctx, key, value, ttl)
 	suite.NoError(err)
 
-	time.Sleep(2 * cleanupInterval)
-
-	var result string
-	err = cache.Get(ctx, key, &result)
+	// Use polling to check for key expiration
+	expired := false
+	maxWaitTime := 10 * cleanupInterval       // Maximum wait time
+	pollingInterval := 100 * time.Millisecond // Polling interval
+	timeout := time.After(maxWaitTime)
+	tick := time.Tick(pollingInterval)
+
+	for !expired {
+		select {
+		case <-timeout:
+			suite.FailNow("Timeout reached, key did not expire as expected")
+			return
+		case <-tick:
+			var result string
+			err = cache.Get(ctx, key, &result)
+			if errors.Is(err, bigcache.ErrEntryNotFound) {
+				expired = true
+			}
+		}
+	}
 
-	suite.True(errors.Is(err, bigcache.ErrEntryNotFound), "Expected 'key not found' error, but got a different error")
+	// Ensure the key has expired as expected
+	suite.True(expired, "Key should have expired")
 }
 
 func (suite *LocalCacheTestSuite) TestCleanupProcessBatchLogic() {
@@ -179,7 +197,7 @@ func (suite *LocalCacheTestSuite) TestCleanupProcessItemDoesNotExpire() {
 // after cleanup, there will be 1 key in the cache and 1 deleted expired key.
 func (suite *LocalCacheTestSuite) TestCleanupProcessMetrics() {
 	cleanupInterval := 1 * time.Second
-	ttl := 10 * time.Millisecond
+	ttl := 25 * time.Millisecond
 
 	// label:count
 	expected := map[string]int{
@@ -226,11 +244,29 @@ func (suite *LocalCacheTestSuite) TestCleanupProcessMetrics() {
 	err = cache.Set(ctx, key+"2", value, neverExpires)
 	suite.NoError(err)
 
-	time.Sleep(2 * cleanupInterval)
-	for k, v := range expected {
-		gotV, ok := got.Load(k)
-		suite.Assert().True(ok)
-
-		suite.Assert().Equal(v, gotV)
+	// Polling to check if metrics have been updated as expected
+	maxWaitTime := 10 * cleanupInterval       // Maximum wait time
+	pollingInterval := 100 * time.Millisecond // Polling interval
+	timeout := time.After(maxWaitTime)
+	tick := time.Tick(pollingInterval)
+
+	for len(expected) > 0 {
+		select {
+		case <-timeout:
+			suite.FailNow("Timeout reached, metrics were not updated as expected")
+			return
+		case <-tick:
+			for k, v := range expected {
+				if gotV, ok := got.Load(k); ok && v == gotV {
+					got.Delete(k)
+					delete(expected, k)
+				}
+			}
+			if len(expected) == 0 {
+				return
+			}
+		}
 	}
+
+	suite.FailNow("Not all expected metrics were updated as expected")
 }

pkg/zrouter/domain/response.go

@@ -9,6 +9,7 @@ import (
 const (
 	ContentTypeHeader          = "Content-Type"
 	ContentTypeApplicationJSON = "application/json; charset=utf-8"
+	ContentTypePlainText       = "text/plain; charset=utf-8"
 	ContentTypeJSON            = "json"
 )
 

pkg/zrouter/handler.go

@@ -14,6 +14,18 @@ func NotFoundHandler(_ Context) (domain.ServiceResponse, error) {
 	return domain.NewErrorNotFound(msg), nil
 }
 
+func ToHandlerFunc(h http.Handler) HandlerFunc {
+	return func(ctx Context) (domain.ServiceResponse, error) {
+		chiCtx, ok := ctx.(*chiContextAdapter)
+		if !ok {
+			return nil, errors.New("context provided is not a *chiContextAdapter")
+		}
+
+		h.ServeHTTP(chiCtx.ctx, chiCtx.req)
+		return nil, nil
+	}
+}
+
 func getChiHandler(handler HandlerFunc) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
 		adaptedContext := &chiContextAdapter{ctx: w, req: r}

pkg/zrouter/zrouter.go

@@ -79,19 +79,23 @@ type ZRouter interface {
 }
 
 type Routes interface {
+	NewSubRouter() ZRouter
 	GET(path string, handler HandlerFunc, middlewares ...zmiddlewares.Middleware) Routes
 	POST(path string, handler HandlerFunc, middlewares ...zmiddlewares.Middleware) Routes
 	PUT(path string, handler HandlerFunc, middlewares ...zmiddlewares.Middleware) Routes
 	PATCH(path string, handler HandlerFunc, middlewares ...zmiddlewares.Middleware) Routes
 	DELETE(path string, handler HandlerFunc, middlewares ...zmiddlewares.Middleware) Routes
+	Handle(pattern string, handler HandlerFunc)
 	Route(method, path string, handler HandlerFunc, middlewares ...zmiddlewares.Middleware) Routes
+	Mount(pattern string, subRouter Routes)
 	Group(prefix string) Routes
 	Use(middlewares ...zmiddlewares.Middleware) Routes
 	NoRoute(handler HandlerFunc)
 	GetRegisteredRoutes() []RegisteredRoute
 	SetDefaultMiddlewares(loggingOptions zmiddlewares.LoggingMiddlewareOptions)
 	GetHandler() http.Handler
 	ServeHTTP(w http.ResponseWriter, req *http.Request)
+	ServeFiles(routePattern string, httpHandler http.Handler)
 }
 
 type zrouter struct {
@@ -132,6 +136,16 @@ func New(metricsServer metrics.TaskMetrics, config *Config) ZRouter {
 	return zr
 }
 
+func (r *zrouter) NewSubRouter() ZRouter {
+	newRouter := &zrouter{
+		router:        chi.NewRouter(),
+		metricsServer: r.metricsServer,
+		config:        r.config,
+	}
+
+	return newRouter
+}
+
 func (r *zrouter) SetDefaultMiddlewares(loggingOptions zmiddlewares.LoggingMiddlewareOptions) {
 	if err := zmiddlewares.RegisterRequestMetrics(r.metricsServer); err != nil {
 		logger.GetLoggerFromContext(context.Background()).Errorf("Error registering metrics %v", err)
@@ -262,6 +276,23 @@ func (r *zrouter) Use(middlewares ...zmiddlewares.Middleware) Routes {
 	return r
 }
 
+func (r *zrouter) Mount(pattern string, subRouter Routes) {
+	sr, ok := subRouter.(*zrouter)
+	if !ok {
+		panic("subRouter is not of expected type *zrouter")
+	}
+
+	r.router.Mount(pattern, sr.router)
+}
+
+func (r *zrouter) Handle(pattern string, handler HandlerFunc) {
+	if handler == nil {
+		panic("handler is mandatory")
+	}
+
+	r.router.Handle(pattern, getChiHandler(handler))
+}
+
 func (r *zrouter) GetRegisteredRoutes() []RegisteredRoute {
 	r.mutex.Lock()
 	defer r.mutex.Unlock()
@@ -279,6 +310,10 @@ func (r *zrouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	r.router.ServeHTTP(w, req)
 }
 
+func (r *zrouter) ServeFiles(routePattern string, httpHandler http.Handler) {
+	r.router.Handle(routePattern, httpHandler)
+}
+
 func LogTopJWTPathMetrics(ctx context.Context, zCache zcache.RemoteCache, updateInterval time.Duration, topN int) {
 	if updateInterval == 0 {
 		updateInterval = defaultUpdateInterval

pkg/zrouter/zrouter_mock.go

@@ -3,6 +3,7 @@ package zrouter
 import (
 	"github.com/stretchr/testify/mock"
 	"github.com/zondax/golem/pkg/zrouter/zmiddlewares"
+	"net/http"
 )
 
 type MockZRouter struct {
@@ -53,3 +54,37 @@ func (m *MockZRouter) Group(prefix string) Routes {
 	args := m.Called(prefix)
 	return args.Get(0).(Routes)
 }
+
+func (m *MockZRouter) Handle(pattern string, handler HandlerFunc) {
+	m.Called(pattern, handler)
+}
+
+func (m *MockZRouter) Mount(pattern string, handler HandlerFunc) {
+	m.Called(pattern, handler)
+}
+
+func (m *MockZRouter) ServeFiles(routePattern string, httpHandler http.Handler) {
+	m.Called(routePattern, httpHandler)
+}
+
+func (m *MockZRouter) NoRoute(handler HandlerFunc) {
+	m.Called(handler)
+}
+
+func (m *MockZRouter) GetRegisteredRoutes() []RegisteredRoute {
+	args := m.Called()
+	return args.Get(0).([]RegisteredRoute)
+}
+
+func (m *MockZRouter) SetDefaultMiddlewares(loggingOptions zmiddlewares.LoggingMiddlewareOptions) {
+	m.Called(loggingOptions)
+}
+
+func (m *MockZRouter) GetHandler() http.Handler {
+	args := m.Called()
+	return args.Get(0).(http.Handler)
+}
+
+func (m *MockZRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	m.Called(w, req)
+}