This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/main by this push:
new 23ea0fadf8 test(pyamber): add unit test coverage for AssignPortHandler
(#7374)
23ea0fadf8 is described below
commit 23ea0fadf85578a45c82c828fb2ce5b9e48bd494
Author: Eugene Gu <[email protected]>
AuthorDate: Thu Aug 6 20:46:50 2026 -0700
test(pyamber): add unit test coverage for AssignPortHandler (#7374)
### What changes were proposed in this PR?
Adds
`amber/src/test/python/core/architecture/handlers/control/test_assign_port_handler.py`
(14 test functions, 16 collected), the first dedicated tests for the
Python worker's `AssignPort` RPC handler. Sibling handlers in the same
package already have test files; this handler's only driver was
`test_main_loop.py`, whose `AssignPortRequest` fixtures pass no
`storage_uris`/`partitionings`, so the per-URI channel registration loop
never executed and the output branch ran only as setup.
Covered behavior:
- Input branch: `add_input_port` receives the port id, the wrapped
`Schema`, and the `storage_uris`/`partitionings` lists verbatim; one
data channel is registered per URI, in URI order, with the reader
identity derived by `get_from_actor_id_for_input_port_storage`, this
worker as the target, and `is_control` false; with no URIs nothing is
registered. Unlike the output branch there is no truthiness guard, so an
empty-string URI still registers a channel — the Scala worker behaves
the same, pinned as a cross-language symmetry.
- Output branch: no URIs and an empty-string URI both yield
`storage_uri_base=None` (the production coordinator really sends `[""]`
for a storage-less output port — `RegionExecutionManager` falls back to
`.getOrElse("")`), a single URI is passed through, and every URI after
the first is dropped (pinned as current behavior).
- Both branches return `EmptyReturn`, wrap `req.schema` into a
value-equal `Schema`, and leave the other manager untouched; an
`AssignPortRequest` built without the `storage_uris` field iterates as
an empty repeated field.
The harness follows `test_start_worker_handler.py`: `asyncio.run` per
test, `SimpleNamespace` context with `MagicMock` managers, no production
code changes.
### Any related issues, documentation, discussions?
Closes #7371.
### How was this PR tested?
85 passed (69 baseline + 16 new). An independent review applied 14
targeted mutations to the handler (branch swap, dropped registration
loop, swapped channel endpoints, `is_control` flip, reordered/emptied
lists, wrong URI index, dropped or added truthiness guards on either
branch, `None` base, `None` return, empty schema in each branch, wrong
port id, fully stubbed body) — every mutation is killed by at least one
test. `black --check` and `ruff check` are clean.
### Was this PR authored or co-authored using generative AI tooling?
Co-authored by: Claude Code (Fable 5)
---
.../handlers/control/test_assign_port_handler.py | 223 +++++++++++++++++++++
1 file changed, 223 insertions(+)
diff --git
a/amber/src/test/python/core/architecture/handlers/control/test_assign_port_handler.py
b/amber/src/test/python/core/architecture/handlers/control/test_assign_port_handler.py
new file mode 100644
index 0000000000..c18bc38057
--- /dev/null
+++
b/amber/src/test/python/core/architecture/handlers/control/test_assign_port_handler.py
@@ -0,0 +1,223 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from core.architecture.handlers.control.assign_port_handler import
AssignPortHandler
+from core.models import Schema
+from core.util.virtual_identity import get_from_actor_id_for_input_port_storage
+from proto.org.apache.texera.amber.core import (
+ ActorVirtualIdentity,
+ ChannelIdentity,
+ PortIdentity,
+)
+from proto.org.apache.texera.amber.engine.architecture.rpc import (
+ AssignPortRequest,
+ EmptyReturn,
+)
+from proto.org.apache.texera.amber.engine.architecture.sendsemantics import (
+ OneToOnePartitioning,
+ Partitioning,
+)
+
+WORKER_ID = "worker-1"
+
+PORT_ID = PortIdentity(id=3, internal=False)
+
+# Two attributes so ordering survives the round trip through the raw map.
+RAW_SCHEMA = {"name": "STRING", "count": "INTEGER"}
+
+
+def _expected_channel(uri: str) -> ChannelIdentity:
+ """The channel the handler must register for one storage URI: derived
+ through the same production util the handler calls, so the assertion pins
+ the real MATERIALIZATION_READER_<uri><worker> identity rather than a
stub."""
+ to_actor_id = ActorVirtualIdentity(WORKER_ID)
+ return ChannelIdentity(
+ get_from_actor_id_for_input_port_storage(uri, to_actor_id),
+ to_actor_id,
+ False,
+ )
+
+
+def _make_request(**overrides) -> AssignPortRequest:
+ fields = dict(
+ port_id=PORT_ID,
+ input=True,
+ schema=RAW_SCHEMA,
+ storage_uris=[],
+ partitionings=[],
+ )
+ fields.update(overrides)
+ return AssignPortRequest(**fields)
+
+
+class TestAssignPortHandler:
+ @pytest.fixture
+ def handler(self):
+ context = SimpleNamespace(
+ input_manager=MagicMock(),
+ output_manager=MagicMock(),
+ worker_id=WORKER_ID,
+ )
+ return AssignPortHandler(context)
+
+ # --- input branch -----------------------------------------------------
+
+ def test_input_without_storage_adds_the_port_without_registering(self,
handler):
+ asyncio.run(handler.assign_port(_make_request()))
+ handler.context.input_manager.add_input_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), [], []
+ )
+ handler.context.input_manager.register_input.assert_not_called()
+
+ def test_input_branch_leaves_the_output_manager_untouched(self, handler):
+ asyncio.run(handler.assign_port(_make_request()))
+ assert handler.context.output_manager.mock_calls == []
+
+ def test_input_passes_storage_uris_and_partitionings_verbatim(self,
handler):
+ uris = ["vfs:///mat/a", "vfs:///mat/b"]
+ partitionings = [
+
Partitioning(one_to_one_partitioning=OneToOnePartitioning(batch_size=1)),
+
Partitioning(one_to_one_partitioning=OneToOnePartitioning(batch_size=2)),
+ ]
+ request = _make_request(storage_uris=uris, partitionings=partitionings)
+ asyncio.run(handler.assign_port(request))
+ handler.context.input_manager.add_input_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), uris, partitionings
+ )
+
+ def test_input_registers_one_channel_per_uri_in_uri_order(self, handler):
+ uris = ["vfs:///mat/a", "vfs:///mat/b"]
+ request = _make_request(storage_uris=uris)
+ asyncio.run(handler.assign_port(request))
+ register_input = handler.context.input_manager.register_input
+ assert register_input.call_count == 2
+ assert [call.kwargs for call in register_input.call_args_list] == [
+ {"channel_id": _expected_channel(uri), "port_id": PORT_ID} for uri
in uris
+ ]
+
+ def test_registered_channel_is_a_data_channel_into_this_worker(self,
handler):
+ uri = "vfs:///mat/a"
+ asyncio.run(handler.assign_port(_make_request(storage_uris=[uri])))
+ register_input = handler.context.input_manager.register_input
+ channel_id = register_input.call_args.kwargs["channel_id"]
+ # The reader identity is derived from the URI plus this worker's own
+ # identity, and the channel carries data — not control — traffic.
+ to_actor_id = ActorVirtualIdentity(WORKER_ID)
+ assert channel_id.to_worker_id == to_actor_id
+ assert channel_id.from_worker_id ==
get_from_actor_id_for_input_port_storage(
+ uri, to_actor_id
+ )
+ assert channel_id.is_control is False
+
+ def test_input_registers_a_channel_even_for_an_empty_string_uri(self,
handler):
+ # Unlike the output branch, the input loop has no truthiness guard: an
+ # empty-string URI still gets a channel derived from "". The Scala
+ # worker behaves the same (URI.create("") succeeds and the channel is
+ # registered), so this pins a cross-language symmetry. The production
+ # coordinator currently sends [] for storage-less input ports
+ # (RegionExecutionManager falls back to List.empty), so this path is
+ # latent wire-format space, not a live message shape.
+ request = _make_request(storage_uris=[""])
+ asyncio.run(handler.assign_port(request))
+ handler.context.input_manager.add_input_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), [""], []
+ )
+ handler.context.input_manager.register_input.assert_called_once_with(
+ channel_id=_expected_channel(""), port_id=PORT_ID
+ )
+
+ def test_input_with_unset_storage_uris_field_iterates_as_empty(self,
handler):
+ # betterproto materializes an unset repeated field as [] on read, so a
+ # request built without storage_uris must behave exactly like one built
+ # with an explicit empty list.
+ request = AssignPortRequest(port_id=PORT_ID, input=True,
schema=RAW_SCHEMA)
+ asyncio.run(handler.assign_port(request))
+ handler.context.input_manager.add_input_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), [], []
+ )
+ handler.context.input_manager.register_input.assert_not_called()
+
+ # --- output branch ----------------------------------------------------
+
+ def test_output_without_storage_uris_passes_a_none_base(self, handler):
+ # The Scala twin calls storageUris.head and would throw on an empty
+ # list; the Python guard checks the length first, so an empty list is
+ # tolerated here even though the coordinator never sends one.
+ asyncio.run(handler.assign_port(_make_request(input=False)))
+ handler.context.output_manager.add_output_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), None
+ )
+
+ def test_output_branch_leaves_the_input_manager_untouched(self, handler):
+ asyncio.run(handler.assign_port(_make_request(input=False)))
+ assert handler.context.input_manager.mock_calls == []
+
+ def test_output_treats_an_empty_string_uri_as_no_storage(self, handler):
+ # The production coordinator really sends [""] for a storage-less
+ # output port (RegionExecutionManager wraps the URI lookup in
+ # .getOrElse("")), so the truthiness guard collapsing "" to None is
+ # live protocol, not defensive coding.
+ request = _make_request(input=False, storage_uris=[""])
+ asyncio.run(handler.assign_port(request))
+ handler.context.output_manager.add_output_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), None
+ )
+
+ def test_output_uses_the_single_uri_as_storage_base(self, handler):
+ request = _make_request(input=False, storage_uris=["vfs:///x"])
+ asyncio.run(handler.assign_port(request))
+ handler.context.output_manager.add_output_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), "vfs:///x"
+ )
+
+ def test_output_silently_drops_every_uri_after_the_first(self, handler):
+ # Only storage_uris[0] is consulted; "vfs:///y" vanishes without an
+ # error or a log line (current behavior).
+ request = _make_request(input=False, storage_uris=["vfs:///x",
"vfs:///y"])
+ asyncio.run(handler.assign_port(request))
+ handler.context.output_manager.add_output_port.assert_called_once_with(
+ PORT_ID, Schema(raw_schema=RAW_SCHEMA), "vfs:///x"
+ )
+
+ # --- both branches ----------------------------------------------------
+
+ @pytest.mark.parametrize("is_input", [True, False])
+ def test_returns_empty_return_from_either_branch(self, handler, is_input):
+ result =
asyncio.run(handler.assign_port(_make_request(input=is_input)))
+ assert result == EmptyReturn()
+
+ @pytest.mark.parametrize("is_input", [True, False])
+ def test_wraps_the_raw_request_schema_into_a_schema_object(self, handler,
is_input):
+ asyncio.run(handler.assign_port(_make_request(input=is_input)))
+ manager = (
+ handler.context.input_manager
+ if is_input
+ else handler.context.output_manager
+ )
+ add_port = manager.add_input_port if is_input else
manager.add_output_port
+ schema = add_port.call_args.args[1]
+ # The manager receives a real Schema built from the request's raw
+ # name -> type map, with attribute order and types preserved.
+ assert isinstance(schema, Schema)
+ assert schema == Schema(raw_schema=RAW_SCHEMA)
+ assert schema.get_attr_names() == ["name", "count"]
