def testGenerateShardArgs(self): num_examples = 37 num_shards = 4 outfiles = [str(i) for i in range(num_shards)] shard_args = gene_expression.generate_shard_args(outfiles, num_examples) starts, ends, fnames = zip(*shard_args) self.assertAllEqual([0, 9, 18, 27], starts) self.assertAllEqual([9, 18, 27, 37], ends) self.assertAllEqual(fnames, outfiles)