@@ -99,18 +99,27 @@ def stream_experiment_data(request, **kwargs):
99
99
start_time = 1571084846.25
100
100
start_time_pb = test_util .timestamp_pb (1571084846250000000 )
101
101
102
+ def outdir_files ():
103
+ # Recursively list `outdir`.
104
+ result = []
105
+ for (dirpath , dirnames , filenames ) in os .walk (outdir ):
106
+ for filename in filenames :
107
+ fullpath = os .path .join (dirpath , filename )
108
+ result .append (os .path .relpath (fullpath , outdir ))
109
+ return result
110
+
102
111
generator = exporter .export (read_time = start_time )
103
112
expected_files = []
104
113
self .assertTrue (os .path .isdir (outdir ))
105
- self .assertCountEqual (expected_files , os . listdir ( outdir ))
114
+ self .assertCountEqual (expected_files , outdir_files ( ))
106
115
mock_api_client .StreamExperiments .assert_not_called ()
107
116
mock_api_client .StreamExperimentData .assert_not_called ()
108
117
109
118
# The first iteration should request the list of experiments and
110
119
# data for one of them.
111
120
self .assertEqual (next (generator ), "123" )
112
- expected_files .append ("scalars_123 .json" )
113
- self .assertCountEqual (expected_files , os . listdir ( outdir ))
121
+ expected_files .append (os . path . join ( "experiment_123" , "scalars .json") )
122
+ self .assertCountEqual (expected_files , outdir_files ( ))
114
123
115
124
expected_eids_request = export_service_pb2 .StreamExperimentsRequest ()
116
125
expected_eids_request .read_timestamp .CopyFrom (start_time_pb )
@@ -131,8 +140,8 @@ def stream_experiment_data(request, **kwargs):
131
140
mock_api_client .StreamExperimentData .reset_mock ()
132
141
self .assertEqual (next (generator ), "456" )
133
142
134
- expected_files .append ("scalars_456 .json" )
135
- self .assertCountEqual (expected_files , os . listdir ( outdir ))
143
+ expected_files .append (os . path . join ( "experiment_456" , "scalars .json") )
144
+ self .assertCountEqual (expected_files , outdir_files ( ))
136
145
mock_api_client .StreamExperiments .assert_not_called ()
137
146
expected_data_request .experiment_id = "456"
138
147
mock_api_client .StreamExperimentData .assert_called_once_with (
@@ -141,12 +150,12 @@ def stream_experiment_data(request, **kwargs):
141
150
142
151
# Again, request data for the next experiment; this experiment ID
143
152
# was in the second response batch in the list of IDs.
144
- expected_files .append ("scalars_789 .json" )
153
+ expected_files .append (os . path . join ( "experiment_789" , "scalars .json") )
145
154
mock_api_client .StreamExperiments .reset_mock ()
146
155
mock_api_client .StreamExperimentData .reset_mock ()
147
156
self .assertEqual (next (generator ), "789" )
148
157
149
- self .assertCountEqual (expected_files , os . listdir ( outdir ))
158
+ self .assertCountEqual (expected_files , outdir_files ( ))
150
159
mock_api_client .StreamExperiments .assert_not_called ()
151
160
expected_data_request .experiment_id = "789"
152
161
mock_api_client .StreamExperimentData .assert_called_once_with (
@@ -158,12 +167,14 @@ def stream_experiment_data(request, **kwargs):
158
167
mock_api_client .StreamExperimentData .reset_mock ()
159
168
self .assertEqual (list (generator ), [])
160
169
161
- self .assertCountEqual (expected_files , os . listdir ( outdir ))
170
+ self .assertCountEqual (expected_files , outdir_files ( ))
162
171
mock_api_client .StreamExperiments .assert_not_called ()
163
172
mock_api_client .StreamExperimentData .assert_not_called ()
164
173
165
- # Spot-check one of the files.
166
- with open (os .path .join (outdir , "scalars_456.json" )) as infile :
174
+ # Spot-check one of the scalar data files.
175
+ with open (
176
+ os .path .join (outdir , "experiment_456" , "scalars.json" )
177
+ ) as infile :
167
178
jsons = [json .loads (line ) for line in infile ]
168
179
self .assertLen (jsons , 4 )
169
180
datum = jsons [2 ]
@@ -309,29 +320,6 @@ def test_rejects_existing_directory(self):
309
320
mock_api_client .StreamExperiments .assert_not_called ()
310
321
mock_api_client .StreamExperimentData .assert_not_called ()
311
322
312
- def test_rejects_existing_file (self ):
313
- mock_api_client = self ._create_mock_api_client ()
314
-
315
- def stream_experiments (request , ** kwargs ):
316
- del request # unused
317
- yield export_service_pb2 .StreamExperimentsResponse (
318
- experiment_ids = ["123" ]
319
- )
320
-
321
- mock_api_client .StreamExperiments = stream_experiments
322
-
323
- outdir = os .path .join (self .get_temp_dir (), "outdir" )
324
- exporter = exporter_lib .TensorBoardExporter (mock_api_client , outdir )
325
- generator = exporter .export ()
326
-
327
- with open (os .path .join (outdir , "scalars_123.json" ), "w" ):
328
- pass
329
-
330
- with self .assertRaises (exporter_lib .OutputFileExistsError ):
331
- next (generator )
332
-
333
- mock_api_client .StreamExperimentData .assert_not_called ()
334
-
335
323
def test_propagates_mkdir_errors (self ):
336
324
mock_api_client = self ._create_mock_api_client ()
337
325
outdir = os .path .join (self .get_temp_dir (), "some_file" , "outdir" )
0 commit comments