@@ -28,12 +28,11 @@ import (
2828 fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
2929)
3030
31- func TestStdinServiceWithEcho (t * testing.T ) {
31+ func mustProcessStdinService (t * testing.T , conf * sspb.Config , im * sspb.InputMessage ) * sspb.OutputMessage {
32+ t .Helper ()
3233 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 ),
3736 })
3837 if err != nil {
3938 t .Fatal (err )
@@ -50,9 +49,7 @@ func TestStdinServiceWithEcho(t *testing.T) {
5049 err = s .ProcessMessage (context .Background (),
5150 & fspb.Message {
5251 MessageType : "StdinServiceInputMessage" ,
53- Data : anypbtest .New (t , & sspb.InputMessage {
54- Args : []string {"-c" , `print("foo bar")` },
55- }),
52+ Data : anypbtest .New (t , im ),
5653 })
5754 if err != nil {
5855 t .Fatal (err )
@@ -62,146 +59,80 @@ func TestStdinServiceWithEcho(t *testing.T) {
6259 select {
6360 case output = <- outChan :
6461 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" )
6663 }
6764
6865 om := & sspb.OutputMessage {}
6966 if err := output .Data .UnmarshalTo (om ); err != nil {
7067 t .Fatal (err )
7168 }
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
7970}
8071
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" ,
9075 }
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" },
9878 }
9979
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 )
12184 }
85+ }
12286
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" ,
12890 }
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" ),
13393 }
13494
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 ) {
13999 t .Fatalf ("unexpected output; got %q, want %q" , om .Stdout , wantStdout )
140100 }
141101}
142102
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 ,
152107 }
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" },
160110 }
111+ om := mustProcessStdinService (t , conf , im )
161112
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 )
179116 }
117+ }
180118
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 ,
186123 }
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" ),
191126 }
127+ om := mustProcessStdinService (t , conf , im )
192128
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 )
201132 }
202133}
203134
204- func TestStdinServiceCancellation (t * testing.T ) {
135+ func TestStdinService_Cancelation (t * testing.T ) {
205136 s , err := Factory (& fspb.ClientServiceConfig {
206137 Name : "SleepService" ,
207138 Config : anypbtest .New (t , & sspb.Config {
0 commit comments