def testWriteSimple(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter( self.test_dir, ioctx, max_file_size=1000, compress_columns=["obs"]) self.assertEqual(len(os.listdir(self.test_dir)), 0) writer.write(SAMPLES) writer.write(SAMPLES) self.assertEqual(len(os.listdir(self.test_dir)), 1)
def testWritePaginate(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter( self.test_dir, ioctx, max_file_size=5000, compress_columns=["obs"]) self.assertEqual(len(os.listdir(self.test_dir)), 0) for _ in range(100): writer.write(SAMPLES) self.assertEqual(len(os.listdir(self.test_dir)), 12)
def testWritePaginate(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter(ioctx, self.test_dir, max_file_size=5000, compress_columns=["obs"]) self.assertEqual(len(os.listdir(self.test_dir)), 0) for _ in range(100): writer.write(SAMPLES) self.assertEqual(len(os.listdir(self.test_dir)), 12)
def testWritePaginate(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter( self.test_dir, ioctx, max_file_size=5000, compress_columns=["obs"]) self.assertEqual(len(os.listdir(self.test_dir)), 0) for _ in range(100): writer.write(SAMPLES) num_files = len(os.listdir(self.test_dir)) # Magic numbers: 2: On travis, it seems to create only 2 files. # 12 or 13: Mac locally. # Reasons: Different compressions, file-size interpretations, # json writers? assert num_files in [2, 12, 13], \ "Expected 12|13 files, but found {} ({})". \ format(num_files, os.listdir(self.test_dir))
def testReadWrite(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter( self.test_dir, ioctx, max_file_size=5000, compress_columns=["obs"]) for i in range(100): writer.write(make_sample_batch(i)) reader = JsonReader(self.test_dir + "/*.json") seen_a = set() seen_o = set() for i in range(1000): batch = reader.next() seen_a.add(batch["actions"][0]) seen_o.add(batch["obs"][0]) self.assertGreater(len(seen_a), 90) self.assertLess(len(seen_a), 101) self.assertGreater(len(seen_o), 90) self.assertLess(len(seen_o), 101)
def test_write_paginate(self): ioctx = IOContext(self.test_dir, {}, 0, None) writer = JsonWriter( self.test_dir, ioctx, max_file_size=5000, compress_columns=["obs"]) self.assertEqual(len(os.listdir(self.test_dir)), 0) for _ in range(100): writer.write(SAMPLES) num_files = len(os.listdir(self.test_dir)) # Pagination can't really be predicted: # On travis, it seems to create only 2 files, but sometimes also # 6, or 7. 12 or 13 usually on a Mac locally. # Reasons: Different compressions, file-size interpretations, # json writers? assert num_files >= 2, \ "Expected >= 2 files, but found {} ({})". \ format(num_files, os.listdir(self.test_dir))