MCPX MCP Gateway

package streams import ( "errors" "fmt" lunar_messages "lunar/engine/messages" stream_config "lunar/engine/streams/config" test_processors "lunar/engine/streams/flow/test-processors" internal_types "lunar/engine/streams/internal-types" lunar_context "lunar/engine/streams/lunar-context" metrics_data "lunar/engine/streams/metrics-data" "lunar/engine/streams/processors" filter_processor "lunar/engine/streams/processors/filter-processor" public_types "lunar/engine/streams/public-types" stream_types "lunar/engine/streams/types" "lunar/engine/utils/environment" context_manager "lunar/toolkit-core/context-manager" "os" "path/filepath" "testing" "time" "github.com/stretchr/testify/require" ) var sharedState = lunar_context.NewMemoryState[[]byte]() func TestMain(m *testing.M) { prevVal := environment.SetProcessorsDirectory(filepath.Join("flow", "test-processors")) context_manager.Get().SetMockClock() // Run the tests code := m.Run() // Clean up if necessary environment.SetProcessorsDirectory(prevVal) // Exit with the code from the tests os.Exit(code) } func setFlowRepDirectory(path string) string { return environment.SetStreamsFlowsDirectory(path) } func revertFlowRepDirectory(path string) { environment.SetStreamsFlowsDirectory(path) } func TestNewStream(t *testing.T) { stream, err := NewStream() require.NoError(t, err, "Failed to create stream") require.NotNil(t, stream, "stream is nil") require.NotNil(t, stream.apiStreams, "APIStreams is nil") require.NotNil(t, stream.filterTree, "filterTree is nil") } func TestRequestBodyFromResponseStream(t *testing.T) { reqBody := "request body" resBody := "response body" seqID := "1234" apiStreamReq := stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeRequest, sharedState) apiStreamReq.SetRequest(stream_types.NewRequest(lunar_messages.OnRequest{ Method: "GET", Scheme: "https", SequenceID: seqID, URL: "maps.googleapis.com/maps/api/geocode/json", Headers: map[string]string{}, RawBody: []byte(reqBody), })) apiStreamReq.StoreRequest() apiStreamRes := stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeResponse, sharedState) apiStreamRes.SetResponse(stream_types.NewResponse(lunar_messages.OnResponse{ Method: "GET", SequenceID: seqID, URL: "maps.googleapis.com/maps/api/geocode/json", Headers: map[string]string{}, RawBody: []byte(resBody), })) req := apiStreamRes.GetRequest() fmt.Printf("Request body: %+v\v", req) fmt.Printf("Request: %s\v", req.GetBody()) require.Equal(t, reqBody, apiStreamRes.GetRequest().GetBody(), "Request body is not correct") require.Equal(t, resBody, apiStreamRes.GetResponse().GetBody(), "Response body is not correct") } func TestExecuteFlows(t *testing.T) { procMng := createTestProcessorManager( t, []string{ "removePII", "readCache", "checkLimit", "generateResponse", "globalStream", "writeCache", "LogAPM", "readXXX", "writeXXX", }, ) stream, err := NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng flowReps := createFlowRepresentation(t, "2-flows-test*") defer revertFlowRepDirectory(setFlowRepDirectory(filepath.Join("flow", "test-cases", "2-flows-test-case"))) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") err = stream.createFlows(flowReps) require.NoError(t, err, "Failed to create flows") apiStream := stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeRequest, sharedState) apiStream.SetResponse(stream_types.NewResponse(lunar_messages.OnResponse{ Status: 200, })) apiStream.SetRequest(stream_types.NewRequest(lunar_messages.OnRequest{ Method: "GET", Scheme: "https", URL: "maps.googleapis.com/maps/api/geocode/json", Headers: map[string]string{}, })) flowActions := &stream_config.StreamActions{ Request: &stream_config.RequestStream{}, Response: &stream_config.ResponseStream{}, } err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") apiStream.SetType(public_types.StreamTypeResponse) err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") // Test for 3 flows procMng = createTestProcessorManager( t, []string{ "removePII", "readCache", "checkLimit", "generateResponse", "globalStream", "writeCache", "LogAPM", "readXXX", "writeXXX", }, ) stream, err = NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng flowReps = createFlowRepresentation(t, "3-flows*") setFlowRepDirectory(filepath.Join("flow", "test-cases", "3-flows-test-case")) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") err = stream.createFlows(flowReps) require.NoError(t, err, "Failed to create flows") apiStream = stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeRequest, sharedState) apiStream.SetResponse(stream_types.NewResponse(lunar_messages.OnResponse{ Status: 200, })) apiStream.SetRequest(stream_types.NewRequest(lunar_messages.OnRequest{ Method: "GET", Scheme: "https", URL: "www.whatever.com/blabla", Headers: map[string]string{}, })) err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") } func TestCreateFlows(t *testing.T) { procMng := createTestProcessorManager( t, []string{ "removePII", "readCache", "checkLimit", "generateResponse", "globalStream", "writeCache", "LogAPM", "readXXX", "writeXXX", }, ) stream, err := NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng flowReps := createFlowRepresentation(t, "2-flows-test*") defer revertFlowRepDirectory(setFlowRepDirectory(filepath.Join("flow", "test-cases", "2-flows-test-case"))) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") err = stream.createFlows(flowReps) require.NoError(t, err, "Failed to create flows") procMng = createTestProcessorManager( t, []string{ "removePII", "readCache", "checkLimit", "generateResponse", "globalStream", "writeCache", "LogAPM", "readXXX", "writeXXX", }, ) stream, err = NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng flowReps = createFlowRepresentation(t, "3-flows*") setFlowRepDirectory(filepath.Join("flow", "test-cases", "3-flows-test-case")) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") err = stream.createFlows(flowReps) require.NoError(t, err, "Failed to create flows") } func TestCreateFlowsWithSameProcessorsName(t *testing.T) { procMng := createTestProcessorManager( t, []string{ "removePII", "readCache", "checkLimit", "generateResponse", "globalStream", "writeCache", "LogAPM", "readCache", "writeCache", }, ) stream, err := NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng flowReps := createFlowRepresentation(t, "2-flows-same-processor-key-test-case") defer revertFlowRepDirectory( setFlowRepDirectory(filepath.Join("flow", "test-cases", "2-flows-same-processor-key-test-case")), ) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") err = stream.createFlows(flowReps) require.NoError(t, err, "Failed to create flows") } func TestEarlyResponseFlow(t *testing.T) { procMng := createTestProcessorManagerWithFactories( t, []string{"readCache", "writeCache", "generateResponse", "LogAPM"}, test_processors.NewMockProcessorUsingCache, test_processors.NewMockProcessor, test_processors.NewMockGenerateResponseProcessor, test_processors.NewMockProcessor, ) stream, err := NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng defer revertFlowRepDirectory(setFlowRepDirectory(filepath.Join("flow", "test-cases", "early-response-test-case"))) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") require.NoError(t, err, "Failed to create flows") contextManager := lunar_context.NewContextManager() globalContext := contextManager.GetGlobalContext() apiStream := stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeRequest, sharedState) apiStream.SetResponse(stream_types.NewResponse(lunar_messages.OnResponse{ Method: "GET", Status: 200, URL: "maps.googleapis.com/maps/api/geocode/json", })) apiStream.SetRequest(stream_types.NewRequest(lunar_messages.OnRequest{ Method: "GET", Scheme: "https", URL: "maps.googleapis.com/maps/api/geocode/json", Headers: map[string]string{}, })) flowActions := &stream_config.StreamActions{ Request: &stream_config.RequestStream{}, Response: &stream_config.ResponseStream{}, } // simulate early response err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") err = globalContext.Set(test_processors.GlobalKeyCacheHit, true) require.NoError(t, err, "Failed to set global context value") err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err := globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"readCache", "generateResponse", "LogAPM"}, execOrder, "Execution order is not correct") // simulate regular execution apiStream.SetType(public_types.StreamTypeRequest) err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") err = globalContext.Set(test_processors.GlobalKeyCacheHit, false) require.NoError(t, err, "Failed to set global context value") err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err = globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"readCache"}, execOrder, "Execution order is not correct") // simulate API provider response apiStream.SetType(public_types.StreamTypeResponse) err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err = globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"writeCache"}, execOrder, "Execution order is not correct") } func TestEarlyResponseFromAnotherFlow(t *testing.T) { procMng := createTestProcessorManager( t, []string{ "removePII", "readCache", "checkLimit", "generateResponse", "globalStream", "writeCache", "LogAPM", "readCache", "writeCache", }, ) stream, err := NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng defer revertFlowRepDirectory(setFlowRepDirectory(filepath.Join("flow", "test-cases", "cross-flow-processor-use"))) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") contextManager := lunar_context.NewContextManager() globalContext := contextManager.GetGlobalContext() apiStream := stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeRequest, sharedState) apiStream.SetResponse(stream_types.NewResponse(lunar_messages.OnResponse{ Status: 200, URL: "maps.googleapis.com/maps/api/geocode/json", })) apiStream.SetRequest(stream_types.NewRequest(lunar_messages.OnRequest{ Method: "GET", Scheme: "https", URL: "maps.googleapis.com/maps/api/geocode/json", Headers: map[string]string{}, })) flowActions := &stream_config.StreamActions{ Request: &stream_config.RequestStream{}, Response: &stream_config.ResponseStream{}, } // simulate early response err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") err = globalContext.Set(test_processors.GlobalKeyCacheHit, true) require.NoError(t, err, "Failed to set global context value") err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err := globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"readCache"}, execOrder, "Execution order is not correct") } func TestFilterProcessorFlow(t *testing.T) { procMng := createTestProcessorManagerWithFactories(t, []string{"Filter", "generateResponse", "LogAPM"}, filter_processor.NewProcessor, test_processors.NewMockGenerateResponseProcessor, test_processors.NewMockProcessor, ) stream, err := NewStream() require.NoError(t, err, "Failed to create stream") stream.processorsManager = procMng _ = createFlowRepresentation(t, "filter*") defer revertFlowRepDirectory(setFlowRepDirectory(filepath.Join("flow", "test-cases", "filter-processor-test-case"))) err = stream.Initialize() require.NoError(t, err, "Failed to create flows") require.NoError(t, err, "Failed to create flows") contextManager := lunar_context.NewContextManager() globalContext := contextManager.GetGlobalContext() apiStream := stream_types.NewAPIStream("APIStreamName", public_types.StreamTypeRequest, sharedState) request := lunar_messages.OnRequest{ Method: "GET", Scheme: "https", URL: "maps.googleapis.com/maps/api/geocode/json", Headers: map[string]string{"x-group": "production"}, } apiStream.SetRequest(stream_types.NewRequest(request)) // execution for production environment err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") flowActions := &stream_config.StreamActions{ Request: &stream_config.RequestStream{}, Response: &stream_config.ResponseStream{}, } err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err := globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"LogAPM"}, execOrder, "Execution order is not correct") // execution for staging environment request.Headers["x-group"] = "staging" apiStream.SetRequest(stream_types.NewRequest(request)) err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err = globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"GenerateResponseTooManyRequests"}, execOrder, "Execution order is not correct") // execution for development environment request.Headers["x-group"] = "development" apiStream.SetRequest(stream_types.NewRequest(request)) err = globalContext.Set(test_processors.GlobalKeyExecutionOrder, []string{}) require.NoError(t, err, "Failed to set global context value") err = stream.ExecuteFlow(apiStream, flowActions) require.NoError(t, err, "Failed to execute flow") execOrder, err = globalContext.Get(test_processors.GlobalKeyExecutionOrder) require.NoError(t, err, "Failed to get global context value") require.Equal(t, []string{"LogAPM"}, execOrder, "Execution order is not correct") } func TestMeasureFlowExecutionTime(t *testing.T) { metrics := metrics_data.NewFlowMetricsData() for i := 0; i < 5; i++ { err := metrics.MeasureFlowExecutionTime("test flow", func() error { time.Sleep(50 * time.Millisecond) // Simulate some work return nil }) require.NoError(t, err) } avgTimeData := metrics.GetAvgFlowExecutionTime() require.Greater(t, avgTimeData.AvgFlowExecutionTime["test flow"], 0.0) // Ensure that the average is within the expected range // Allowing some delta (±10ms) due to possible variations in execution time require.InDelta(t, 50, avgTimeData.AvgFlowExecutionTime["test flow"], 10) // Call the measured function with an error to check that it doesn't affect the average err := metrics.MeasureFlowExecutionTime("test flow", func() error { return errors.New("test error") }) require.Error(t, err) require.Equal(t, "test error", err.Error()) // average execution time should be still the same after the error avgTimeAfterError := metrics.GetAvgFlowExecutionTime() require.InDelta( t, avgTimeData.AvgFlowExecutionTime["test flow"], avgTimeAfterError.AvgFlowExecutionTime["test flow"], 5e-2, // allow ±0.05 "AvgFlowExecutionTime should match within 50ms", ) } func createTestProcessorManager(t *testing.T, processorNames []string) *processors.ProcessorManager { return createTestProcessorManagerWithFactories(t, processorNames, test_processors.NewMockProcessor) } func createTestProcessorManagerWithFactories( t *testing.T, processorNames []string, factories ...processors.ProcessorFactory, ) *processors.ProcessorManager { processorMng := processors.NewProcessorManager(nil) for i, procName := range processorNames { factory := factories[0] if len(factories) > 1 { factory = factories[i] } processorMng.SetFactory(procName, factory) } err := processorMng.Init() require.NoError(t, err) return processorMng } func createFlowRepresentation(t *testing.T, testCase string) map[string]internal_types.FlowRepI { pattern := filepath.Join("flow", "test-cases", testCase, "*.yaml") files, fileErr := filepath.Glob(pattern) require.NoError(t, fileErr, "Failed to find YAML files") flowReps := make(map[string]internal_types.FlowRepI) for _, file := range files { t.Run(filepath.Base(file), func(t *testing.T) { flowRep, err := stream_config.ReadStreamFlowConfig(file) require.NoError(t, err, "Failed to read YAML file") // TODO: We should do it more strict that we load the processor keys for key, proc := range flowRep.Processors { proc.Key = key flowRep.Processors[key] = proc } flowReps[flowRep.Name] = flowRep }) } return flowReps }