Is there a way to set the number of segments required when creating a test data set for an integration test?

Not sure if this is exactly what you are looking for, but in one of my unit test cases I created a table over two segments in the following way:

@BeforeClass
public void setUp()
    throws Exception {
  FileUtils.deleteDirectory(INDEX_DIR);

  List<GenericRow> records1 = new ArrayList<>(NUM_RECORDS);
  records1.add(createRecord(120, 200.50F, "albert1", "albert", 1643666769000L));
  records1.add(createRecord(250, 32.50F, "martian1", "mouse", 1643666728000L));
  records1.add(createRecord(310, -44.50F, "martian2", "mouse", 1643666432000L));
  records1.add(createRecord(340, 11.50F, "donald1", "duck", 1643666726000L));
  records1.add(createRecord(110, 16, "goofy1", "goofy", 1643667762000L));
  records1.add(createRecord(150, 12, "goofy2", "goofy", 1643667762000L));
  records1.add(createRecord(100, -28, "daffy1", "daffy", 1643667092000L));
  records1.add(createRecord(120, -16, "pluto1", "dwag", 1643666712000L));
  records1.add(createRecord(120, -16, "zebra1", "zookeeper", 1643666712000L));
  records1.add(createRecord(220, -16, "zebra2", "zookeeper", 1643666712000L));
  createSegment(records1, SEGMENT_NAME_LEFT);
  ImmutableSegment immutableSegment1 =
      ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME_LEFT), ReadMode.mmap);

  List<GenericRow> records2 = new ArrayList<>(NUM_RECORDS);
  records2.add(createRecord(150, 10.50F, "alice1", "wonderland", 1650069985000L));
  records2.add(createRecord(200, 1.50F, "albert2", "albert", 1650050085000L));
  records2.add(createRecord(32, 10.0F, "mickey1", "mouse", 1650040085000L));
  records2.add(createRecord(-40, 250F, "minney2", "mouse", 1650043085000L));
  records2.add(createRecord(10, 4.50F, "donald2", "duck", 1650011085000L));
  records2.add(createRecord(5, 7.50F, "goofy3", "duck", 1650010085000L));
  records2.add(createRecord(5, 4.50F, "daffy2", "duck", 1650045085000L));
  records2.add(createRecord(10, 46.0F, "daffy3", "duck", 1650032085000L));
  records2.add(createRecord(20, 20.5F, "goofy4", "goofy", 1650011085000L));
  records2.add(createRecord(-20, 2.5F, "pluto2", "dwag", 1650052285000L));
  createSegment(records2, SEGMENT_NAME_RIGHT);
  ImmutableSegment immutableSegment2 =
      ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME_RIGHT), ReadMode.mmap);

  _indexSegment = null;
  _indexSegments = Arrays.asList(immutableSegment1, immutableSegment2);
}