schemas/graph/cosmo_test.go

package graph

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gopkg.in/yaml.v3"

	"gitea.unbound.se/unboundsoftware/schemas/graph/model"
)

// MockCommandExecutor implements CommandExecutor for testing
type MockCommandExecutor struct {
	// CallCount tracks how many times Execute was called
	CallCount int
	// LastArgs stores the arguments from the last call
	LastArgs []string
	// Error can be set to simulate command failures
	Error error
}

// Execute mocks the wgc command by generating a realistic config.json file
func (m *MockCommandExecutor) Execute(name string, args ...string) ([]byte, error) {
	m.CallCount++
	m.LastArgs = append([]string{name}, args...)

	if m.Error != nil {
		return nil, m.Error
	}

	// Parse the input file to understand what subgraphs we're composing
	var inputFile, outputFile string
	for i, arg := range args {
		if arg == "--input" && i+1 < len(args) {
			inputFile = args[i+1]
		}
		if arg == "--out" && i+1 < len(args) {
			outputFile = args[i+1]
		}
	}

	if inputFile == "" || outputFile == "" {
		return nil, fmt.Errorf("missing required arguments")
	}

	// Read the input YAML to get subgraph information
	inputData, err := os.ReadFile(inputFile)
	if err != nil {
		return nil, fmt.Errorf("failed to read input file: %w", err)
	}

	var input struct {
		Version   int `yaml:"version"`
		Subgraphs []struct {
			Name         string                 `yaml:"name"`
			RoutingURL   string                 `yaml:"routing_url,omitempty"`
			Schema       map[string]string      `yaml:"schema"`
			Subscription map[string]interface{} `yaml:"subscription,omitempty"`
		} `yaml:"subgraphs"`
	}

	if err := yaml.Unmarshal(inputData, &input); err != nil {
		return nil, fmt.Errorf("failed to parse input YAML: %w", err)
	}

	// Generate a realistic Cosmo Router config based on the input
	config := map[string]interface{}{
		"version": "mock-version-uuid",
		"subgraphs": func() []map[string]interface{} {
			subgraphs := make([]map[string]interface{}, len(input.Subgraphs))
			for i, sg := range input.Subgraphs {
				subgraph := map[string]interface{}{
					"id":   fmt.Sprintf("mock-id-%d", i),
					"name": sg.Name,
				}
				if sg.RoutingURL != "" {
					subgraph["routingUrl"] = sg.RoutingURL
				}
				subgraphs[i] = subgraph
			}
			return subgraphs
		}(),
		"engineConfig": map[string]interface{}{
			"graphqlSchema": generateMockSchema(input.Subgraphs),
			"datasourceConfigurations": func() []map[string]interface{} {
				dsConfigs := make([]map[string]interface{}, len(input.Subgraphs))
				for i, sg := range input.Subgraphs {
					// Read SDL from file
					sdl := ""
					if schemaFile, ok := sg.Schema["file"]; ok {
						if sdlData, err := os.ReadFile(schemaFile); err == nil {
							sdl = string(sdlData)
						}
					}

					dsConfig := map[string]interface{}{
						"id":   fmt.Sprintf("datasource-%d", i),
						"kind": "GRAPHQL",
						"customGraphql": map[string]interface{}{
							"federation": map[string]interface{}{
								"enabled":    true,
								"serviceSdl": sdl,
							},
							"subscription": func() map[string]interface{} {
								if len(sg.Subscription) > 0 {
									return map[string]interface{}{
										"enabled": true,
										"url": map[string]interface{}{
											"staticVariableContent": sg.Subscription["url"],
										},
										"protocol":             sg.Subscription["protocol"],
										"websocketSubprotocol": sg.Subscription["websocket_subprotocol"],
									}
								}
								return map[string]interface{}{
									"enabled": false,
								}
							}(),
						},
					}
					dsConfigs[i] = dsConfig
				}
				return dsConfigs
			}(),
		},
	}

	// Write the config to the output file
	configJSON, err := json.Marshal(config)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal config: %w", err)
	}

	if err := os.WriteFile(outputFile, configJSON, 0o644); err != nil {
		return nil, fmt.Errorf("failed to write output file: %w", err)
	}

	return []byte("Success"), nil
}

// generateMockSchema creates a simple merged schema from subgraphs
func generateMockSchema(subgraphs []struct {
	Name         string                 `yaml:"name"`
	RoutingURL   string                 `yaml:"routing_url,omitempty"`
	Schema       map[string]string      `yaml:"schema"`
	Subscription map[string]interface{} `yaml:"subscription,omitempty"`
},
) string {
	schema := strings.Builder{}
	schema.WriteString("schema {\n  query: Query\n")

	// Check if any subgraph has subscriptions
	hasSubscriptions := false
	for _, sg := range subgraphs {
		if len(sg.Subscription) > 0 {
			hasSubscriptions = true
			break
		}
	}

	if hasSubscriptions {
		schema.WriteString("  subscription: Subscription\n")
	}
	schema.WriteString("}\n\n")

	// Add types by reading SDL files
	for _, sg := range subgraphs {
		if schemaFile, ok := sg.Schema["file"]; ok {
			if sdlData, err := os.ReadFile(schemaFile); err == nil {
				schema.WriteString(string(sdlData))
				schema.WriteString("\n")
			}
		}
	}

	return schema.String()
}

func TestGenerateCosmoRouterConfig(t *testing.T) {
	tests := []struct {
		name      string
		subGraphs []*model.SubGraph
		wantErr   bool
		validate  func(t *testing.T, config string)
	}{
		{
			name: "single subgraph with all fields",
			subGraphs: []*model.SubGraph{
				{
					Service: "test-service",
					URL:     stringPtr("http://localhost:4001/query"),
					WsURL:   stringPtr("ws://localhost:4001/query"),
					Sdl:     "type Query { test: String }",
				},
			},
			wantErr: false,
			validate: func(t *testing.T, config string) {
				var result map[string]interface{}
				err := json.Unmarshal([]byte(config), &result)
				require.NoError(t, err, "Config should be valid JSON")

				// Version is a UUID string from wgc
				version, ok := result["version"].(string)
				require.True(t, ok, "Version should be a string")
				assert.NotEmpty(t, version, "Version should not be empty")

				subgraphs, ok := result["subgraphs"].([]interface{})
				require.True(t, ok, "subgraphs should be an array")
				require.Len(t, subgraphs, 1, "Should have 1 subgraph")

				sg := subgraphs[0].(map[string]interface{})
				assert.Equal(t, "test-service", sg["name"])
				assert.Equal(t, "http://localhost:4001/query", sg["routingUrl"])

				// Check that datasource configurations include subscription settings
				engineConfig, ok := result["engineConfig"].(map[string]interface{})
				require.True(t, ok, "Should have engineConfig")

				dsConfigs, ok := engineConfig["datasourceConfigurations"].([]interface{})
				require.True(t, ok && len(dsConfigs) > 0, "Should have datasource configurations")

				ds := dsConfigs[0].(map[string]interface{})
				customGraphql, ok := ds["customGraphql"].(map[string]interface{})
				require.True(t, ok, "Should have customGraphql config")

				subscription, ok := customGraphql["subscription"].(map[string]interface{})
				require.True(t, ok, "Should have subscription config")
				assert.True(t, subscription["enabled"].(bool), "Subscription should be enabled")

				subUrl, ok := subscription["url"].(map[string]interface{})
				require.True(t, ok, "Should have subscription URL")
				assert.Equal(t, "ws://localhost:4001/query", subUrl["staticVariableContent"])
			},
		},
		{
			name: "multiple subgraphs",
			subGraphs: []*model.SubGraph{
				{
					Service: "service-1",
					URL:     stringPtr("http://localhost:4001/query"),
					Sdl:     "type Query { field1: String }",
				},
				{
					Service: "service-2",
					URL:     stringPtr("http://localhost:4002/query"),
					Sdl:     "type Query { field2: String }",
				},
				{
					Service: "service-3",
					URL:     stringPtr("http://localhost:4003/query"),
					WsURL:   stringPtr("ws://localhost:4003/query"),
					Sdl:     "type Subscription { updates: String }",
				},
			},
			wantErr: false,
			validate: func(t *testing.T, config string) {
				var result map[string]interface{}
				err := json.Unmarshal([]byte(config), &result)
				require.NoError(t, err)

				subgraphs := result["subgraphs"].([]interface{})
				assert.Len(t, subgraphs, 3, "Should have 3 subgraphs")

				// Check service names
				sg1 := subgraphs[0].(map[string]interface{})
				assert.Equal(t, "service-1", sg1["name"])

				sg3 := subgraphs[2].(map[string]interface{})
				assert.Equal(t, "service-3", sg3["name"])

				// Check that datasource configurations include subscription for service-3
				engineConfig, ok := result["engineConfig"].(map[string]interface{})
				require.True(t, ok, "Should have engineConfig")

				dsConfigs, ok := engineConfig["datasourceConfigurations"].([]interface{})
				require.True(t, ok && len(dsConfigs) == 3, "Should have 3 datasource configurations")

				// Find service-3's datasource config (should have subscription enabled)
				ds3 := dsConfigs[2].(map[string]interface{})
				customGraphql, ok := ds3["customGraphql"].(map[string]interface{})
				require.True(t, ok, "Service-3 should have customGraphql config")

				subscription, ok := customGraphql["subscription"].(map[string]interface{})
				require.True(t, ok, "Service-3 should have subscription config")
				assert.True(t, subscription["enabled"].(bool), "Service-3 subscription should be enabled")
			},
		},
		{
			name: "subgraph with no URL",
			subGraphs: []*model.SubGraph{
				{
					Service: "test-service",
					URL:     nil,
					WsURL:   nil,
					Sdl:     "type Query { test: String }",
				},
			},
			wantErr: false,
			validate: func(t *testing.T, config string) {
				var result map[string]interface{}
				err := json.Unmarshal([]byte(config), &result)
				require.NoError(t, err)

				subgraphs := result["subgraphs"].([]interface{})
				sg := subgraphs[0].(map[string]interface{})

				// Should not have routing URL when URL is nil
				_, hasRoutingURL := sg["routingUrl"]
				assert.False(t, hasRoutingURL, "Should not have routingUrl when URL is nil")

				// Check datasource configurations don't have subscription enabled
				engineConfig, ok := result["engineConfig"].(map[string]interface{})
				require.True(t, ok, "Should have engineConfig")

				dsConfigs, ok := engineConfig["datasourceConfigurations"].([]interface{})
				require.True(t, ok && len(dsConfigs) > 0, "Should have datasource configurations")

				ds := dsConfigs[0].(map[string]interface{})
				customGraphql, ok := ds["customGraphql"].(map[string]interface{})
				require.True(t, ok, "Should have customGraphql config")

				subscription, ok := customGraphql["subscription"].(map[string]interface{})
				if ok {
					// wgc always enables subscription but URL should be empty when WsURL is nil
					subUrl, hasUrl := subscription["url"].(map[string]interface{})
					if hasUrl {
						_, hasStaticContent := subUrl["staticVariableContent"]
						assert.False(t, hasStaticContent, "Subscription URL should be empty when WsURL is nil")
					}
				}
			},
		},
		{
			name:      "empty subgraphs",
			subGraphs: []*model.SubGraph{},
			wantErr:   true,
			validate:  nil,
		},
		{
			name:      "nil subgraphs",
			subGraphs: nil,
			wantErr:   true,
			validate:  nil,
		},
		{
			name: "complex SDL with multiple types",
			subGraphs: []*model.SubGraph{
				{
					Service: "complex-service",
					URL:     stringPtr("http://localhost:4001/query"),
					Sdl: `
						type Query {
							user(id: ID!): User
							users: [User!]!
						}

						type User {
							id: ID!
							name: String!
							email: String!
						}
					`,
				},
			},
			wantErr: false,
			validate: func(t *testing.T, config string) {
				var result map[string]interface{}
				err := json.Unmarshal([]byte(config), &result)
				require.NoError(t, err)

				// Check the composed graphqlSchema contains the types
				engineConfig, ok := result["engineConfig"].(map[string]interface{})
				require.True(t, ok, "Should have engineConfig")

				graphqlSchema, ok := engineConfig["graphqlSchema"].(string)
				require.True(t, ok, "Should have graphqlSchema")

				assert.Contains(t, graphqlSchema, "Query", "Schema should contain Query type")
				assert.Contains(t, graphqlSchema, "User", "Schema should contain User type")

				// Check datasource has the original SDL
				dsConfigs, ok := engineConfig["datasourceConfigurations"].([]interface{})
				require.True(t, ok && len(dsConfigs) > 0, "Should have datasource configurations")

				ds := dsConfigs[0].(map[string]interface{})
				customGraphql, ok := ds["customGraphql"].(map[string]interface{})
				require.True(t, ok, "Should have customGraphql config")

				federation, ok := customGraphql["federation"].(map[string]interface{})
				require.True(t, ok, "Should have federation config")

				serviceSdl, ok := federation["serviceSdl"].(string)
				require.True(t, ok, "Should have serviceSdl")
				assert.Contains(t, serviceSdl, "email: String!", "SDL should contain email field")
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Use mock executor for all tests
			mockExecutor := &MockCommandExecutor{}
			config, err := GenerateCosmoRouterConfigWithExecutor(tt.subGraphs, mockExecutor)

			if tt.wantErr {
				assert.Error(t, err)
				// Verify executor was not called for error cases
				if len(tt.subGraphs) == 0 {
					assert.Equal(t, 0, mockExecutor.CallCount, "Should not call executor for empty subgraphs")
				}
				return
			}

			require.NoError(t, err)
			assert.NotEmpty(t, config, "Config should not be empty")

			// Verify executor was called correctly
			assert.Equal(t, 1, mockExecutor.CallCount, "Should call executor once")
			assert.Equal(t, "wgc", mockExecutor.LastArgs[0], "Should call wgc command")
			assert.Contains(t, mockExecutor.LastArgs, "router", "Should include 'router' arg")
			assert.Contains(t, mockExecutor.LastArgs, "compose", "Should include 'compose' arg")

			if tt.validate != nil {
				tt.validate(t, config)
			}
		})
	}
}

// TestGenerateCosmoRouterConfig_MockError tests error handling with mock executor
func TestGenerateCosmoRouterConfig_MockError(t *testing.T) {
	subGraphs := []*model.SubGraph{
		{
			Service: "test-service",
			URL:     stringPtr("http://localhost:4001/query"),
			Sdl:     "type Query { test: String }",
		},
	}

	// Create a mock executor that returns an error
	mockExecutor := &MockCommandExecutor{
		Error: fmt.Errorf("simulated wgc failure"),
	}

	config, err := GenerateCosmoRouterConfigWithExecutor(subGraphs, mockExecutor)

	// Verify error is propagated
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "wgc router compose failed")
	assert.Contains(t, err.Error(), "simulated wgc failure")
	assert.Empty(t, config)

	// Verify executor was called
	assert.Equal(t, 1, mockExecutor.CallCount, "Should have attempted to call executor")
}

// SlowMockExecutor simulates a slow wgc command for concurrency testing.
type SlowMockExecutor struct {
	MockCommandExecutor
	delay      time.Duration
	mu         sync.Mutex
	concurrent atomic.Int32
	maxSeen    atomic.Int32
}

func (m *SlowMockExecutor) Execute(name string, args ...string) ([]byte, error) {
	cur := m.concurrent.Add(1)
	// Track the maximum concurrent executions observed.
	for {
		old := m.maxSeen.Load()
		if cur <= old || m.maxSeen.CompareAndSwap(old, cur) {
			break
		}
	}
	defer m.concurrent.Add(-1)

	time.Sleep(m.delay)

	m.mu.Lock()
	defer m.mu.Unlock()
	return m.MockCommandExecutor.Execute(name, args...)
}

func TestCosmoGenerator_ConcurrencyLimit(t *testing.T) {
	executor := &SlowMockExecutor{delay: 100 * time.Millisecond}
	gen := NewCosmoGenerator(executor, 5*time.Second)

	subGraphs := []*model.SubGraph{
		{
			Service: "svc",
			URL:     stringPtr("http://localhost:4001/query"),
			Sdl:     "type Query { hello: String }",
		},
	}

	var wg sync.WaitGroup
	for range 5 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, _ = gen.Generate(context.Background(), subGraphs)
		}()
	}
	wg.Wait()

	assert.Equal(t, int32(1), executor.maxSeen.Load(),
		"at most 1 wgc process should run concurrently")
}

func TestCosmoGenerator_Timeout(t *testing.T) {
	// Executor that takes longer than the timeout.
	executor := &SlowMockExecutor{delay: 500 * time.Millisecond}
	gen := NewCosmoGenerator(executor, 50*time.Millisecond)

	subGraphs := []*model.SubGraph{
		{
			Service: "svc",
			URL:     stringPtr("http://localhost:4001/query"),
			Sdl:     "type Query { hello: String }",
		},
	}

	// First call: occupies the semaphore for 500ms.
	go func() {
		_, _ = gen.Generate(context.Background(), subGraphs)
	}()

	// Give the first goroutine time to acquire the semaphore.
	time.Sleep(20 * time.Millisecond)

	// Second call: should timeout waiting for the semaphore.
	_, err := gen.Generate(context.Background(), subGraphs)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "acquire cosmo generator")
}

func TestCosmoGenerator_ContextCancellation(t *testing.T) {
	executor := &SlowMockExecutor{delay: 500 * time.Millisecond}
	gen := NewCosmoGenerator(executor, 5*time.Second)

	subGraphs := []*model.SubGraph{
		{
			Service: "svc",
			URL:     stringPtr("http://localhost:4001/query"),
			Sdl:     "type Query { hello: String }",
		},
	}

	// First call: occupies the semaphore.
	go func() {
		_, _ = gen.Generate(context.Background(), subGraphs)
	}()

	time.Sleep(20 * time.Millisecond)

	// Second call with an already-cancelled context.
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err := gen.Generate(ctx, subGraphs)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "acquire cosmo generator")
}

// Helper function for tests
func stringPtr(s string) *string {
	return &s
}