go-sdk/schema_test.go at main · streamdal/go-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
package streamdal

import (
	"context"
	"sync"
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	"github.com/streamdal/streamdal/libs/protos/build/go/protos"
	"github.com/streamdal/streamdal/libs/protos/build/go/protos/steps"

	"github.com/streamdal/go-sdk/logger/loggerfakes"
	"github.com/streamdal/go-sdk/server/serverfakes"
)

var _ = Describe("Schema", func() {

	var s *Streamdal
	var aud *protos.Audience

	BeforeEach(func() {
		s = &Streamdal{
			schemasMtx: &sync.RWMutex{},
			schemas:    make(map[string]*protos.Schema),
			config: &Config{
				Logger: &loggerfakes.FakeLogger{},
			},
		}

		aud = &protos.Audience{
			ServiceName:   "test-service",
			OperationType: protos.OperationType_OPERATION_TYPE_PRODUCER,
			OperationName: "test-operation",
			ComponentName: "kafka",
		}
	})

	Context("getSchema", func() {
		It("returns empty when no schema found", func() {
			schema := s.getSchema(context.Background(), nil)
			Expect(schema).To(Equal([]byte(``)))
		})

		It("returns schema when found", func() {
			s.schemas[audToStr(aud)] = &protos.Schema{
				JsonSchema: []byte(`{"type": "string"}`),
			}

			schema := s.getSchema(context.Background(), aud)
			Expect(schema).To(Equal([]byte(`{"type": "string"}`)))
		})
	})

	Context("setSchema", func() {
		It("sets schema", func() {
			s.setSchema(context.Background(), aud, []byte(`{"type": "string"}`))
			schema := s.getSchema(context.Background(), aud)
			Expect(schema).To(Equal([]byte(`{"type": "string"}`)))
		})
	})

	Context("handleSchema", func() {
		It("returns false when no schema step", func() {
			fakeClient := &serverfakes.FakeIServerClient{}
			s.serverClient = fakeClient

			step := &protos.PipelineStep{
				Step: nil,
			}

			got := s.handleSchema(context.Background(), aud, step, nil)

			Expect(got).To(BeFalse())
		})

		It("returns false on bad wasm exit", func() {
			fakeClient := &serverfakes.FakeIServerClient{}
			s.serverClient = fakeClient

			step := &protos.PipelineStep{
				Step: &protos.PipelineStep_InferSchema{
					InferSchema: &steps.InferSchemaStep{
						CurrentSchema: []byte(`{}`),
					},
				},
			}

			wasmResp := &protos.WASMResponse{
				ExitCode: protos.WASMExitCode_WASM_EXIT_CODE_FALSE,
			}

			got := s.handleSchema(context.Background(), aud, step, wasmResp)

			Expect(got).To(BeFalse())
		})

		It("sends to server when an updated schema is generated", func() {
			fakeClient := &serverfakes.FakeIServerClient{}
			s.serverClient = fakeClient

			step := &protos.PipelineStep{
				Step: &protos.PipelineStep_InferSchema{
					InferSchema: &steps.InferSchemaStep{
						CurrentSchema: []byte(`{}`),
					},
				},
			}

			wasmResp := &protos.WASMResponse{
				ExitCode:   protos.WASMExitCode_WASM_EXIT_CODE_TRUE,
				OutputStep: []byte(`{"type": "string"}`),
			}

			got := s.handleSchema(context.Background(), aud, step, wasmResp)

			// Give time for goroutine to run
			time.Sleep(time.Millisecond * 500)

			Expect(got).To(BeTrue())
			Expect(fakeClient.SendSchemaCallCount()).To(Equal(1))
		})
	})
})