| #!/usr/bin/env vpython3 |
| # Copyright 2025 The LUCI Authors. All rights reserved. |
| # Use of this source code is governed under the Apache License, Version 2.0 |
| # that can be found in the LICENSE file. |
| |
| from __future__ import annotations |
| |
| from typing import cast |
| |
| from unittest.mock import MagicMock |
| |
| from google.protobuf.timestamp_pb2 import Timestamp |
| |
| import test_env |
| import turboci_test_helper |
| |
| from google.protobuf.message import Message |
| |
| from recipe_engine.internal.turboci.ids import type_urls |
| |
| from google.protobuf.struct_pb2 import Struct, Value as StructValue |
| from google.protobuf.proto_json import parse |
| |
| from PB.turboci.graph.ids.v1 import identifier |
| from PB.turboci.graph.orchestrator.v1.check_kind import CheckKind |
| from PB.turboci.graph.orchestrator.v1.check_state import CheckState |
| from PB.turboci.graph.orchestrator.v1.edge import Edge |
| from PB.turboci.graph.orchestrator.v1.query import Query |
| from PB.turboci.graph.orchestrator.v1.query_nodes_request import QueryNodesRequest |
| from PB.turboci.graph.orchestrator.v1.revision import Revision |
| from PB.turboci.graph.orchestrator.v1.type_info import TypeInfo |
| from PB.turboci.graph.orchestrator.v1.type_set import TypeSet |
| from PB.turboci.graph.orchestrator.v1.value_ref import ValueRef |
| from PB.turboci.graph.orchestrator.v1.value_write import ValueWrite |
| from PB.turboci.graph.orchestrator.v1.write_nodes_request import WriteNodesRequest |
| |
| from recipe_engine.internal.turboci import common |
| from recipe_engine.turboci import (write_nodes, reason, check, dep_group, |
| check_id, query_nodes, make_query, |
| get_option, get_results) |
| |
| |
| def _mkStruct(d: dict) -> Struct: |
| return cast(Struct, parse(Struct, d)) |
| |
| |
| def _mkValue(msg: Message, realm: str | None = None) -> ValueWrite: |
| ret = ValueWrite() |
| if realm: |
| ret.realm = realm |
| ret.data.Pack(msg, deterministic=True) |
| return ret |
| |
| |
| class TestCheckID(test_env.RecipeEngineUnitTest): |
| def test_ok(self): |
| self.assertEqual(check_id('fleem'), identifier.Check(id='fleem')) |
| self.assertEqual(check_id('fleem', in_workplan='123'), identifier.Check( |
| work_plan=identifier.WorkPlan(id='123'), |
| id='fleem', |
| )) |
| self.assertEqual(check_id('fleem', in_workplan='L321'), identifier.Check( |
| work_plan=identifier.WorkPlan(id='321'), |
| id='fleem', |
| )) |
| |
| def test_fail(self): |
| with self.assertRaisesRegex(ValueError, 'must not contain'): |
| check_id('not:cool') |
| |
| with self.assertRaises(ValueError) as ex: |
| check_id('just fine', in_workplan='bad news') |
| self.assertRegex(ex.exception.__notes__[0], 'in_workplan: id must be parsable') |
| |
| |
| class TestReason(test_env.RecipeEngineUnitTest): |
| def test_ok(self): |
| r = reason("some string", _mkStruct({'a': 'b'})) |
| |
| self.assertEqual( |
| r, |
| WriteNodesRequest.Reason( |
| message="some string", |
| details=[_mkValue(_mkStruct({'a': 'b'}))], |
| )) |
| |
| |
| class TestDepGroup(test_env.RecipeEngineUnitTest): |
| def test_ok(self): |
| dg = dep_group( |
| 'stuff', |
| identifier.Identifier(check=identifier.Check(id="things")), |
| identifier.Check(id="more"), |
| dep_group('nested-1', 'nested-2', in_workplan='L123456'), |
| threshold=3, |
| ) |
| self.assertEqual( |
| dg, |
| WriteNodesRequest.DependencyGroup( |
| edges=[ |
| Edge(check=Edge.Check(identifier=check_id('stuff'))), |
| Edge(check=Edge.Check(identifier=check_id('things'))), |
| Edge(check=Edge.Check(identifier=check_id('more'))), |
| ], |
| groups=[ |
| WriteNodesRequest.DependencyGroup(edges=[ |
| Edge(check=Edge.Check(identifier=check_id( |
| 'nested-1', in_workplan='123456'))), |
| Edge(check=Edge.Check(identifier=check_id( |
| 'nested-2', in_workplan='123456'))), |
| ]), |
| ], |
| threshold=3, |
| )) |
| |
| |
| class TestCheck(test_env.RecipeEngineUnitTest): |
| def test_ok(self): |
| chk = check( |
| 'the check id', |
| kind='CHECK_KIND_TEST', |
| state='CHECK_STATE_PLANNED', |
| options=[_mkStruct({'a': 'b'})], |
| deps=dep_group( |
| "stuff", |
| "things", |
| ), |
| results=[_mkStruct({'cool': ['result']})], |
| finalize_results=True, |
| in_workplan='321', |
| realm='project/check/realm', |
| realm_options=[ |
| ('project/check/option/realm', |
| StructValue(string_value='realm_option')), |
| ], |
| realm_results=[ |
| ('project/check/result/realm', |
| StructValue(string_value='realm_result')), |
| ], |
| ) |
| |
| self.assertEqual( |
| chk, |
| WriteNodesRequest.CheckWrite( |
| identifier=identifier.Check( |
| work_plan=identifier.WorkPlan(id='321'), |
| id='the check id', |
| ), |
| realm='project/check/realm', |
| kind=CheckKind.CHECK_KIND_TEST, |
| options=[ |
| _mkValue(_mkStruct({'a': 'b'})), |
| _mkValue( |
| StructValue(string_value='realm_option'), |
| realm='project/check/option/realm'), |
| ], |
| dependencies=WriteNodesRequest.DependencyGroup( |
| edges=[ |
| Edge(check=Edge.Check(identifier=check_id('stuff'))), |
| Edge(check=Edge.Check(identifier=check_id('things'))), |
| ],), |
| results=[ |
| _mkValue(_mkStruct({'cool': ['result']})), |
| _mkValue( |
| StructValue(string_value='realm_result'), |
| realm='project/check/result/realm'), |
| ], |
| finalize_results=True, |
| state=CheckState.CHECK_STATE_PLANNED, |
| )) |
| |
| |
| class TestWriteNodes(test_env.RecipeEngineUnitTest): |
| |
| def setUp(self): |
| self.m = MagicMock() |
| common.CLIENT = self.m |
| return super().setUp() |
| |
| def tearDown(self): |
| common.CLIENT = None |
| return super().tearDown() |
| |
| def test_write_nodes(self): |
| # User writes: |
| write_nodes( |
| check( |
| "someid", |
| kind='CHECK_KIND_BUILD', |
| options=[ |
| _mkStruct({"cool_opt": [1, 2, 3]}), |
| ]), |
| reason("I feel like it", _mkStruct({"hello": "world"})), |
| ) |
| |
| # Raw API call to common.CLIENT. |
| self.m.WriteNodes.assert_called_once_with( |
| WriteNodesRequest( |
| reason=WriteNodesRequest.Reason( |
| message="I feel like it", |
| details=[_mkValue(_mkStruct({'hello': 'world'}))], |
| ), |
| checks=[ |
| WriteNodesRequest.CheckWrite( |
| identifier=identifier.Check(id="someid"), |
| kind=CheckKind.CHECK_KIND_BUILD, |
| options=[_mkValue(_mkStruct({'cool_opt': [1, 2, 3]}))], |
| ), |
| ], |
| )) |
| |
| def test_query_nodes(self): |
| query_nodes( |
| make_query( |
| node_set=[check_id("bob")], |
| ), make_query( |
| Query.SelectChecks.Predicate(kind='CHECK_KIND_TEST'), |
| Query.CollectChecks(options=True), |
| ), |
| version=QueryNodesRequest.VersionRestriction( |
| require=Revision(ts=Timestamp(seconds=1234, nanos=5678)), |
| )) |
| |
| self.m.QueryNodes.assert_called_once_with( |
| QueryNodesRequest( |
| type_info=TypeInfo(wanted=TypeSet()), |
| query=[ |
| Query( |
| nodes_by_id=Query.NodesByID(nodes=[ |
| identifier.Identifier(check=identifier.Check(id="bob")) |
| ]),), |
| Query( |
| nodes_in_workplan=identifier.WorkPlan(), |
| select_checks=Query.SelectChecks(predicates=[ |
| Query.SelectChecks.Predicate(kind='CHECK_KIND_TEST'), |
| ]), |
| collect_checks=Query.CollectChecks(options=True,), |
| ), |
| ], |
| version=QueryNodesRequest.VersionRestriction( |
| require=Revision(ts=Timestamp(seconds=1234, nanos=5678)),), |
| )) |
| |
| |
| class TestGetOptionsResults(turboci_test_helper.TestBaseClass): |
| |
| def test_get_option(self): |
| self.write_nodes( |
| check( |
| 'a', |
| kind='CHECK_KIND_ANALYSIS', |
| options=[_mkStruct({ |
| 'hello': [1, 2, 3, 4], |
| })])) |
| |
| returned_check = self.read_checks( |
| 'a', |
| collect=Query.CollectChecks(options=True), |
| types=list(type_urls(Struct)))[0] |
| |
| self.assertEqual( |
| get_option(Struct, returned_check), { |
| 'hello': [1, 2, 3, 4], |
| }) |
| |
| def test_get_result(self): |
| self.write_nodes( |
| check( |
| 'a', |
| kind='CHECK_KIND_ANALYSIS', |
| state='CHECK_STATE_WAITING', |
| results=[_mkStruct({ |
| 'hello': [1, 2, 3, 4], |
| })])) |
| |
| returned_check = self.read_checks( |
| 'a', |
| collect=Query.CollectChecks(result_data=True), |
| types=list(type_urls(Struct)))[0] |
| |
| self.assertEqual( |
| get_results(Struct, returned_check), [{ |
| 'hello': [1, 2, 3, 4], |
| }]) |
| |
| |
| if __name__ == '__main__': |
| test_env.main() |
