code: def test_aggregate_public_partitions_drop_non_public(self):
      # Arrange
      aggregator_params = pipeline_dp.AggregateParams(
          noise_kind=pipeline_dp.NoiseKind.GAUSSIAN,
          metrics=[
              agg.Metrics.COUNT,
              agg.Metrics.SUM,
              agg.Metrics.PRIVACY_ID_COUNT,
          ],
          min_value=0,
          max_value=1,
          max_partitions_contributed=1,
          max_contributions_per_partition=1,
      )
  
      # Set an arbitrary budget, we are not interested in the DP outputs, only
      # the partition keys.
      budget_accountant = NaiveBudgetAccountant(
          total_epsilon=1, total_delta=1e-10
      )
  
      # Input collection has 10 elements, such that each privacy id
      # contributes 1 time and each partition has 1 element.
      col = list(range(10))
      data_extractor = pipeline_dp.DataExtractors(
          privacy_id_extractor=lambda x: x,
          partition_extractor=lambda x: f"pk{x}",
          value_extractor=lambda x: x,
      )
  
      engine = pipeline_dp.DPEngine(
          budget_accountant=budget_accountant,
          backend=pipeline_dp.LocalBackend(),
      )
  
      col = engine.aggregate(
          col=col,
          params=aggregator_params,
          data_extractors=data_extractor,
          public_partitions=["pk0", "pk1", "pk10"],
      )
      budget_accountant.compute_budgets()
  
      col = list(col)
      partition_keys = [x[0] for x in col]
      # Assert
  
      # Only public partitions (0, 1, 2) should be kept and the rest of the
      # partitions should be dropped.
      self.assertEqual(["pk0", "pk1", "pk10"], partition_keys)
  
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 28 for calling the DpEngineTest obj's underlying member method code with example usage: obj.test_aggregate_public_partitions_drop_non_public() without return types