github.com/Jeffail/benthos/v3@v3.65.0/public/service/environment_test.go (about)

     1  package service_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Jeffail/benthos/v3/public/bloblang"
    10  	"github.com/Jeffail/benthos/v3/public/service"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  func walkForSummaries(fn func(func(name string, config *service.ConfigView))) map[string]string {
    16  	summaries := map[string]string{}
    17  	fn(func(name string, config *service.ConfigView) {
    18  		summaries[name] = config.Summary()
    19  	})
    20  	return summaries
    21  }
    22  
    23  func TestEnvironmentAdjustments(t *testing.T) {
    24  	envOne := service.NewEnvironment()
    25  	envTwo := envOne.Clone()
    26  
    27  	assert.NoError(t, envOne.RegisterCache(
    28  		"one_cache", service.NewConfigSpec().Summary("cache one"),
    29  		func(conf *service.ParsedConfig, mgr *service.Resources) (service.Cache, error) {
    30  			return nil, errors.New("cache one err")
    31  		},
    32  	))
    33  	assert.NoError(t, envOne.RegisterInput(
    34  		"one_input", service.NewConfigSpec().Summary("input one"),
    35  		func(conf *service.ParsedConfig, mgr *service.Resources) (service.Input, error) {
    36  			return nil, errors.New("input one err")
    37  		},
    38  	))
    39  	assert.NoError(t, envOne.RegisterOutput(
    40  		"one_output", service.NewConfigSpec().Summary("output one"),
    41  		func(conf *service.ParsedConfig, mgr *service.Resources) (service.Output, int, error) {
    42  			return nil, 0, errors.New("output one err")
    43  		},
    44  	))
    45  	assert.NoError(t, envOne.RegisterProcessor(
    46  		"one_processor", service.NewConfigSpec().Summary("processor one"),
    47  		func(conf *service.ParsedConfig, mgr *service.Resources) (service.Processor, error) {
    48  			return nil, errors.New("processor one err")
    49  		},
    50  	))
    51  	assert.NoError(t, envOne.RegisterRateLimit(
    52  		"one_rate_limit", service.NewConfigSpec().Summary("rate limit one"),
    53  		func(conf *service.ParsedConfig, mgr *service.Resources) (service.RateLimit, error) {
    54  			return nil, errors.New("rate limit one err")
    55  		},
    56  	))
    57  
    58  	assert.Equal(t, "cache one", walkForSummaries(envOne.WalkCaches)["one_cache"])
    59  	assert.Equal(t, "input one", walkForSummaries(envOne.WalkInputs)["one_input"])
    60  	assert.Equal(t, "output one", walkForSummaries(envOne.WalkOutputs)["one_output"])
    61  	assert.Equal(t, "processor one", walkForSummaries(envOne.WalkProcessors)["one_processor"])
    62  	assert.Equal(t, "rate limit one", walkForSummaries(envOne.WalkRateLimits)["one_rate_limit"])
    63  
    64  	assert.NotContains(t, walkForSummaries(envTwo.WalkCaches), "one_cache")
    65  	assert.NotContains(t, walkForSummaries(envTwo.WalkInputs), "one_input")
    66  	assert.NotContains(t, walkForSummaries(envTwo.WalkOutputs), "one_output")
    67  	assert.NotContains(t, walkForSummaries(envTwo.WalkProcessors), "one_processor")
    68  	assert.NotContains(t, walkForSummaries(envTwo.WalkRateLimits), "one_rate_limit")
    69  
    70  	testConfig := `
    71  input:
    72    one_input: {}
    73  pipeline:
    74    processors:
    75      - one_processor: {}
    76  output:
    77    one_output: {}
    78  cache_resources:
    79    - label: foocache
    80      one_cache: {}
    81  rate_limit_resources:
    82    - label: foorl
    83      one_rate_limit: {}
    84  `
    85  
    86  	assert.NoError(t, envOne.NewStreamBuilder().SetYAML(testConfig))
    87  	assert.Error(t, envTwo.NewStreamBuilder().SetYAML(testConfig))
    88  }
    89  
    90  func TestEnvironmentBloblangIsolation(t *testing.T) {
    91  	bEnv := bloblang.NewEnvironment().WithoutFunctions("now")
    92  	require.NoError(t, bEnv.RegisterFunctionV2("meow", bloblang.NewPluginSpec(), func(args *bloblang.ParsedParams) (bloblang.Function, error) {
    93  		return func() (interface{}, error) {
    94  			return "meow", nil
    95  		}, nil
    96  	}))
    97  
    98  	envOne := service.NewEnvironment()
    99  	envOne.UseBloblangEnvironment(bEnv)
   100  
   101  	badConfig := `
   102  pipeline:
   103    processors:
   104      - bloblang: 'root = now()'
   105  `
   106  
   107  	goodConfig := `
   108  pipeline:
   109    processors:
   110      - bloblang: 'root = meow()'
   111  
   112  output:
   113    drop: {}
   114  
   115  logger:
   116    level: OFF
   117  `
   118  
   119  	assert.Error(t, envOne.NewStreamBuilder().SetYAML(badConfig))
   120  
   121  	strmBuilder := envOne.NewStreamBuilder()
   122  	require.NoError(t, strmBuilder.SetYAML(goodConfig))
   123  
   124  	var received []string
   125  	require.NoError(t, strmBuilder.AddConsumerFunc(func(c context.Context, m *service.Message) error {
   126  		b, err := m.AsBytes()
   127  		if err != nil {
   128  			return err
   129  		}
   130  		received = append(received, string(b))
   131  		return nil
   132  	}))
   133  
   134  	pFn, err := strmBuilder.AddProducerFunc()
   135  	require.NoError(t, err)
   136  
   137  	strm, err := strmBuilder.Build()
   138  	require.NoError(t, err)
   139  
   140  	go func() {
   141  		require.NoError(t, strm.Run(context.Background()))
   142  	}()
   143  
   144  	require.NoError(t, pFn(context.Background(), service.NewMessage([]byte("hello world"))))
   145  
   146  	require.NoError(t, strm.StopWithin(time.Second))
   147  	assert.Equal(t, []string{"meow"}, received)
   148  }