code: def test_select_partitions(self):
      # This test is probabilistic, but the parameters were chosen to ensure
      # the test has passed at least 10000 runs.
  
      # Arrange
      params = pipeline_dp.SelectPartitionsParams(
          max_partitions_contributed=1
      )
  
      budget_accountant = NaiveBudgetAccountant(
          total_epsilon=1, total_delta=1e-5
      )
  
      # Generate dataset as a list of (user, partition_key) tuples.
      # There partitions are generated to reflect several scenarios.
  
      # A partition with sufficient amount of users.
      col = [(u, "pk-many-contribs") for u in range(25)]
  
      # A partition with many contributions, but only a few unique users.
      col += [
          (100 + u // 10, "pk-many-contribs-few-users")
          for u in range(30)
      ]
  
      # A partition with few contributions.
      col += [(200 + u, "pk-few-contribs") for u in range(3)]
  
      # Generating 30 partitions, each with the same group of 25 users
      # 25 users is sufficient to keep the partition, but because of
      # contribution bounding, much less users per partition will be kept.
      for i in range(30):
          col += [
              (500 + u, f"few-contribs-after-bound{i}")
              for u in range(25)
          ]
  
      col = list(col)
      data_extractor = pipeline_dp.DataExtractors(
          privacy_id_extractor=lambda x: x[0],
          partition_extractor=lambda x: x[1],
      )
  
      engine = pipeline_dp.DPEngine(
          budget_accountant=budget_accountant,
          backend=pipeline_dp.LocalBackend(),
      )
  
      col = engine.select_partitions(
          col=col, params=params, data_extractors=data_extractor
      )
      budget_accountant.compute_budgets()
  
      col = list(col)
  
      # Assert
      # Only one partition is retained, the one that has many unique _after_
      # applying the "max_partitions_contributed" bound is retained.
      self.assertEqual(["pk-many-contribs"], col)
  
imports: import sys
  import unittest
  from typing import List, Optional, Tuple
  from unittest.mock import patch
  import numpy
  import apache_beam
  import apache_beam.testing.test_pipeline
  import apache_beam.testing.util
  import pydp.algorithms.partition_selection
  from absl.testing import absltest
  from absl.testing import parameterized
  import pipeline_dp
  from pipeline_dp import aggregate_params
  from pipeline_dp import budget_accounting
  from pipeline_dp import dp_computations
  from pipeline_dp.budget_accounting import NaiveBudgetAccountant
  from pipeline_dp.pipeline_backend import PipelineBackend
  from pipeline_dp.pipeline_functions import collect_to_container
  from pipeline_dp.report_generator import ReportGenerator
label_desc: Use image node_id 26 for calling the DpEngineTest obj's underlying member method code with example usage: obj.test_select_partitions() without return types