@@ -28,12 +28,11 @@ import (
28
28
fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
29
29
)
30
30
31
- func TestStdinServiceWithEcho (t * testing.T ) {
31
+ func mustProcessStdinService (t * testing.T , conf * sspb.Config , im * sspb.InputMessage ) * sspb.OutputMessage {
32
+ t .Helper ()
32
33
s , err := Factory (& fspb.ClientServiceConfig {
33
- Name : "EchoService" ,
34
- Config : anypbtest .New (t , & sspb.Config {
35
- Cmd : "python" ,
36
- }),
34
+ Name : "TestService" ,
35
+ Config : anypbtest .New (t , conf ),
37
36
})
38
37
if err != nil {
39
38
t .Fatal (err )
@@ -50,9 +49,7 @@ func TestStdinServiceWithEcho(t *testing.T) {
50
49
err = s .ProcessMessage (context .Background (),
51
50
& fspb.Message {
52
51
MessageType : "StdinServiceInputMessage" ,
53
- Data : anypbtest .New (t , & sspb.InputMessage {
54
- Args : []string {"-c" , `print("foo bar")` },
55
- }),
52
+ Data : anypbtest .New (t , im ),
56
53
})
57
54
if err != nil {
58
55
t .Fatal (err )
@@ -62,146 +59,80 @@ func TestStdinServiceWithEcho(t *testing.T) {
62
59
select {
63
60
case output = <- outChan :
64
61
default :
65
- t .Fatal (". ProcessMessage (/bin/echo foo bar ) expected to produce message, but none found" )
62
+ t .Fatal ("ProcessMessage( ) expected to produce message, but none found" )
66
63
}
67
64
68
65
om := & sspb.OutputMessage {}
69
66
if err := output .Data .UnmarshalTo (om ); err != nil {
70
67
t .Fatal (err )
71
68
}
72
-
73
- wantStdout := []byte ("foo bar\n " )
74
- wantStdoutWin := []byte ("foo bar\r \n " )
75
- if ! bytes .Equal (om .Stdout , wantStdout ) &&
76
- ! bytes .Equal (om .Stdout , wantStdoutWin ) {
77
- t .Fatalf ("unexpected output; got %q, want %q" , om .Stdout , wantStdout )
78
- }
69
+ return om
79
70
}
80
71
81
- func TestStdinServiceWithCat (t * testing.T ) {
82
- s , err := Factory (& fspb.ClientServiceConfig {
83
- Name : "CatService" ,
84
- Config : anypbtest .New (t , & sspb.Config {
85
- Cmd : "python" ,
86
- }),
87
- })
88
- if err != nil {
89
- t .Fatal (err )
72
+ func TestStdinService_AcceptsArgs (t * testing.T ) {
73
+ conf := & sspb.Config {
74
+ Cmd : "echo" ,
90
75
}
91
-
92
- outChan := make (chan * fspb.Message , 1 )
93
- err = s .Start (& clitesting.MockServiceContext {
94
- OutChan : outChan ,
95
- })
96
- if err != nil {
97
- t .Fatal (err )
76
+ im := & sspb.InputMessage {
77
+ Args : []string {"foo bar" },
98
78
}
99
79
100
- err = s .ProcessMessage (context .Background (),
101
- & fspb.Message {
102
- MessageType : "StdinServiceInputMessage" ,
103
- Data : anypbtest .New (t , & sspb.InputMessage {
104
- Args : []string {"-c" , `
105
- try:
106
- my_input = raw_input # Python2 compat
107
- except NameError:
108
- my_input = input
109
-
110
- try:
111
- while True:
112
- print(my_input())
113
- except EOFError:
114
- pass
115
- ` },
116
- Input : []byte ("foo bar" ),
117
- }),
118
- })
119
- if err != nil {
120
- t .Fatalf ("s.ProcessMessage(...) = %q, want success" , err )
80
+ om := mustProcessStdinService (t , conf , im )
81
+ wantStdout := []byte ("foo bar\n " )
82
+ if ! bytes .Equal (om .Stdout , wantStdout ) {
83
+ t .Fatalf ("unexpected output; got %q, want %q" , om .Stdout , wantStdout )
121
84
}
85
+ }
122
86
123
- var output * fspb.Message
124
- select {
125
- case output = <- outChan :
126
- default :
127
- t .Fatal (".ProcessMessage (/bin/cat <<< 'foo bar') expected to produce message, but none found" )
87
+ func TestStdinService_AcceptsStdin (t * testing.T ) {
88
+ conf := & sspb.Config {
89
+ Cmd : "cat" ,
128
90
}
129
-
130
- om := & sspb.OutputMessage {}
131
- if err := output .Data .UnmarshalTo (om ); err != nil {
132
- t .Fatal (err )
91
+ im := & sspb.InputMessage {
92
+ Input : []byte ("foo bar" ),
133
93
}
134
94
135
- wantStdout := [] byte ( "foo bar \n " )
136
- wantStdoutWin := [] byte ( "foo bar \r \n " )
137
- if ! bytes . Equal ( om . Stdout , wantStdout ) &&
138
- ! bytes .Equal (om .Stdout , wantStdoutWin ) {
95
+ om := mustProcessStdinService ( t , conf , im )
96
+
97
+ wantStdout := [] byte ( "foo bar" )
98
+ if ! bytes .Equal (om .Stdout , wantStdout ) {
139
99
t .Fatalf ("unexpected output; got %q, want %q" , om .Stdout , wantStdout )
140
100
}
141
101
}
142
102
143
- func TestStdinServiceReportsResourceUsage (t * testing.T ) {
144
- s , err := Factory (& fspb.ClientServiceConfig {
145
- Name : "BashService" ,
146
- Config : anypbtest .New (t , & sspb.Config {
147
- Cmd : "python" ,
148
- }),
149
- })
150
- if err != nil {
151
- t .Fatal (err )
103
+ func TestStdinService_RejectsArgs (t * testing.T ) {
104
+ conf := & sspb.Config {
105
+ Cmd : "echo" ,
106
+ RejectArgs : true ,
152
107
}
153
-
154
- outChan := make (chan * fspb.Message , 1 )
155
- err = s .Start (& clitesting.MockServiceContext {
156
- OutChan : outChan ,
157
- })
158
- if err != nil {
159
- t .Fatal (err )
108
+ im := & sspb.InputMessage {
109
+ Args : []string {"don't print this" },
160
110
}
111
+ om := mustProcessStdinService (t , conf , im )
161
112
162
- err = s .ProcessMessage (context .Background (),
163
- & fspb.Message {
164
- MessageType : "StdinServiceInputMessage" ,
165
- Data : anypbtest .New (t , & sspb.InputMessage {
166
- // Generate some system (os.listdir) and user (everything else) execution time...
167
- Args : []string {"-c" , `
168
- import os
169
- import time
170
-
171
- t0 = time.time()
172
- while time.time() - t0 < 1.:
173
- os.listdir(".")
174
- ` },
175
- }),
176
- })
177
- if err != nil {
178
- t .Fatal (err )
113
+ wantStdout := []byte ("\n " )
114
+ if ! bytes .Equal (om .Stdout , wantStdout ) {
115
+ t .Fatalf ("unexpected output; got %q, want %q" , om .Stdout , wantStdout )
179
116
}
117
+ }
180
118
181
- var output * fspb.Message
182
- select {
183
- case output = <- outChan :
184
- default :
185
- t .Fatal (".ProcessMessage (/bin/bash ...) expected to produce message, but none found" )
119
+ func TestStdinService_RejectsStdin (t * testing.T ) {
120
+ conf := & sspb.Config {
121
+ Cmd : "cat" ,
122
+ RejectStdin : true ,
186
123
}
187
-
188
- om := & sspb.OutputMessage {}
189
- if err := output .Data .UnmarshalTo (om ); err != nil {
190
- t .Fatal (err )
124
+ im := & sspb.InputMessage {
125
+ Input : []byte ("don't print this" ),
191
126
}
127
+ om := mustProcessStdinService (t , conf , im )
192
128
193
- // We don't test for ResourceUsage.MeanResidentMemory because memory is currently not being
194
- // queried after the process has terminated. It's only queried right after launching the command
195
- // in which case it can be recorded as "0" which would be indistinguishable from it not being set
196
- // at all, resulting in a flaky test case. The fact that the other resource usage metrics have
197
- // been set here is good enough for now.
198
-
199
- if om .Timestamp .Seconds <= 0 {
200
- t .Fatalf ("unexpected output; StdinServiceOutputMessage.timestamp.seconds not set: %q" , om )
129
+ wantStdout := []byte ("" )
130
+ if ! bytes .Equal (om .Stdout , wantStdout ) {
131
+ t .Fatalf ("unexpected output; got %q, want %q" , om .Stdout , wantStdout )
201
132
}
202
133
}
203
134
204
- func TestStdinServiceCancellation (t * testing.T ) {
135
+ func TestStdinService_Cancelation (t * testing.T ) {
205
136
s , err := Factory (& fspb.ClientServiceConfig {
206
137
Name : "SleepService" ,
207
138
Config : anypbtest .New (t , & sspb.Config {
0 commit comments