Skip to content

Commit cef2b55

Browse files
tpoliawtpoliaw
authored
feat: Support mocking devices from dodal modules (#1276)
1 parent 1ab51e8 commit cef2b55

8 files changed

Lines changed: 264 additions & 111 deletions

File tree

‎helm/blueapi/config_schema.json‎

Lines changed: 92 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -64,23 +64,89 @@
6464
"title": "CORSConfig",
6565
"type": "object"
6666
},
67+
"DeviceSource": {
68+
"additionalProperties": false,
69+
"properties": {
70+
"module": {
71+
"description": "Module to be imported",
72+
"title": "Module",
73+
"type": "string"
74+
},
75+
"kind": {
76+
"const": "deviceFunctions",
77+
"default": "deviceFunctions",
78+
"title": "Kind",
79+
"type": "string"
80+
}
81+
},
82+
"required": [
83+
"module"
84+
],
85+
"title": "DeviceSource",
86+
"type": "object"
87+
},
88+
"DodalSource": {
89+
"additionalProperties": false,
90+
"properties": {
91+
"module": {
92+
"description": "Module to be imported",
93+
"title": "Module",
94+
"type": "string"
95+
},
96+
"kind": {
97+
"const": "dodal",
98+
"default": "dodal",
99+
"title": "Kind",
100+
"type": "string"
101+
},
102+
"mock": {
103+
"default": false,
104+
"description": "If true, ophyd_async device connections are mocked",
105+
"title": "Mock",
106+
"type": "boolean"
107+
}
108+
},
109+
"required": [
110+
"module"
111+
],
112+
"title": "DodalSource",
113+
"type": "object"
114+
},
67115
"EnvironmentConfig": {
68116
"additionalProperties": false,
69117
"description": "Config for the RunEngine environment",
70118
"properties": {
71119
"sources": {
72120
"default": [
73121
{
74-
"kind": "planFunctions",
75-
"module": "dodal.plans"
122+
"module": "dodal.plans",
123+
"kind": "planFunctions"
76124
},
77125
{
78-
"kind": "planFunctions",
79-
"module": "dodal.plan_stubs.wrapped"
126+
"module": "dodal.plan_stubs.wrapped",
127+
"kind": "planFunctions"
80128
}
81129
],
82130
"items": {
83-
"$ref": "#/$defs/Source"
131+
"discriminator": {
132+
"mapping": {
133+
"deviceFunctions": "#/$defs/DeviceSource",
134+
"dodal": "#/$defs/DodalSource",
135+
"planFunctions": "#/$defs/PlanSource"
136+
},
137+
"propertyName": "kind"
138+
},
139+
"oneOf": [
140+
{
141+
"$ref": "#/$defs/PlanSource"
142+
},
143+
{
144+
"$ref": "#/$defs/DeviceSource"
145+
},
146+
{
147+
"$ref": "#/$defs/DodalSource"
148+
}
149+
]
84150
},
85151
"title": "Sources",
86152
"type": "array"
@@ -216,6 +282,27 @@
216282
"title": "OIDCConfig",
217283
"type": "object"
218284
},
285+
"PlanSource": {
286+
"additionalProperties": false,
287+
"properties": {
288+
"module": {
289+
"description": "Module to be imported",
290+
"title": "Module",
291+
"type": "string"
292+
},
293+
"kind": {
294+
"const": "planFunctions",
295+
"default": "planFunctions",
296+
"title": "Kind",
297+
"type": "string"
298+
}
299+
},
300+
"required": [
301+
"module"
302+
],
303+
"title": "PlanSource",
304+
"type": "object"
305+
},
219306
"RestConfig": {
220307
"additionalProperties": false,
221308
"properties": {
@@ -296,41 +383,6 @@
296383
"title": "ScratchRepository",
297384
"type": "object"
298385
},
299-
"Source": {
300-
"additionalProperties": false,
301-
"properties": {
302-
"kind": {
303-
"$ref": "#/$defs/SourceKind"
304-
},
305-
"module": {
306-
"anyOf": [
307-
{
308-
"format": "path",
309-
"type": "string"
310-
},
311-
{
312-
"type": "string"
313-
}
314-
],
315-
"title": "Module"
316-
}
317-
},
318-
"required": [
319-
"kind",
320-
"module"
321-
],
322-
"title": "Source",
323-
"type": "object"
324-
},
325-
"SourceKind": {
326-
"enum": [
327-
"planFunctions",
328-
"deviceFunctions",
329-
"dodal"
330-
],
331-
"title": "SourceKind",
332-
"type": "string"
333-
},
334386
"StompConfig": {
335387
"additionalProperties": false,
336388
"description": "Config for connecting to stomp broker",

‎helm/blueapi/values.schema.json‎

Lines changed: 77 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -553,6 +553,52 @@
553553
},
554554
"additionalProperties": false
555555
},
556+
"DeviceSource": {
557+
"title": "DeviceSource",
558+
"type": "object",
559+
"required": [
560+
"module"
561+
],
562+
"properties": {
563+
"kind": {
564+
"title": "Kind",
565+
"default": "deviceFunctions",
566+
"const": "deviceFunctions"
567+
},
568+
"module": {
569+
"title": "Module",
570+
"description": "Module to be imported",
571+
"type": "string"
572+
}
573+
},
574+
"additionalProperties": false
575+
},
576+
"DodalSource": {
577+
"title": "DodalSource",
578+
"type": "object",
579+
"required": [
580+
"module"
581+
],
582+
"properties": {
583+
"kind": {
584+
"title": "Kind",
585+
"default": "dodal",
586+
"const": "dodal"
587+
},
588+
"mock": {
589+
"title": "Mock",
590+
"description": "If true, ophyd_async device connections are mocked",
591+
"default": false,
592+
"type": "boolean"
593+
},
594+
"module": {
595+
"title": "Module",
596+
"description": "Module to be imported",
597+
"type": "string"
598+
}
599+
},
600+
"additionalProperties": false
601+
},
556602
"EnvironmentConfig": {
557603
"title": "EnvironmentConfig",
558604
"description": "Config for the RunEngine environment",
@@ -585,7 +631,17 @@
585631
],
586632
"type": "array",
587633
"items": {
588-
"$ref": "#/$defs/Source"
634+
"oneOf": [
635+
{
636+
"$ref": "#/$defs/PlanSource"
637+
},
638+
{
639+
"$ref": "#/$defs/DeviceSource"
640+
},
641+
{
642+
"$ref": "#/$defs/DodalSource"
643+
}
644+
]
589645
}
590646
}
591647
},
@@ -702,6 +758,26 @@
702758
},
703759
"additionalProperties": false
704760
},
761+
"PlanSource": {
762+
"title": "PlanSource",
763+
"type": "object",
764+
"required": [
765+
"module"
766+
],
767+
"properties": {
768+
"kind": {
769+
"title": "Kind",
770+
"default": "planFunctions",
771+
"const": "planFunctions"
772+
},
773+
"module": {
774+
"title": "Module",
775+
"description": "Module to be imported",
776+
"type": "string"
777+
}
778+
},
779+
"additionalProperties": false
780+
},
705781
"RestConfig": {
706782
"title": "RestConfig",
707783
"type": "object",
@@ -778,40 +854,6 @@
778854
},
779855
"additionalProperties": false
780856
},
781-
"Source": {
782-
"title": "Source",
783-
"type": "object",
784-
"required": [
785-
"kind",
786-
"module"
787-
],
788-
"properties": {
789-
"kind": {
790-
"$ref": "#/$defs/SourceKind"
791-
},
792-
"module": {
793-
"title": "Module",
794-
"anyOf": [
795-
{
796-
"type": "string"
797-
},
798-
{
799-
"type": "string"
800-
}
801-
]
802-
}
803-
},
804-
"additionalProperties": false
805-
},
806-
"SourceKind": {
807-
"title": "SourceKind",
808-
"type": "string",
809-
"enum": [
810-
"planFunctions",
811-
"deviceFunctions",
812-
"dodal"
813-
]
814-
},
815857
"StompConfig": {
816858
"title": "StompConfig",
817859
"description": "Config for connecting to stomp broker",

‎src/blueapi/config.py‎

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from functools import cached_property
77
from pathlib import Path
88
from string import Template
9-
from typing import Any, Generic, Literal, TypeVar, cast
9+
from typing import Annotated, Any, Generic, Literal, TypeVar, cast
1010

1111
import requests
1212
import yaml
@@ -50,8 +50,26 @@ class SourceKind(str, Enum):
5050

5151

5252
class Source(BlueapiBaseModel):
53-
kind: SourceKind
54-
module: Path | str
53+
module: str = Field(description="Module to be imported")
54+
55+
56+
class PlanSource(Source):
57+
kind: Literal[SourceKind.PLAN_FUNCTIONS] = Field(
58+
SourceKind.PLAN_FUNCTIONS, init=False
59+
)
60+
61+
62+
class DeviceSource(Source):
63+
kind: Literal[SourceKind.DEVICE_FUNCTIONS] = Field(
64+
SourceKind.DEVICE_FUNCTIONS, init=False
65+
)
66+
67+
68+
class DodalSource(Source):
69+
kind: Literal[SourceKind.DODAL] = Field(SourceKind.DODAL, init=False)
70+
mock: bool = Field(
71+
description="If true, ophyd_async device connections are mocked", default=False
72+
)
5573

5674

5775
class TcpUrl(AnyUrl):
@@ -101,9 +119,14 @@ class EnvironmentConfig(BlueapiBaseModel):
101119
Config for the RunEngine environment
102120
"""
103121

104-
sources: list[Source] = [
105-
Source(kind=SourceKind.PLAN_FUNCTIONS, module="dodal.plans"),
106-
Source(kind=SourceKind.PLAN_FUNCTIONS, module="dodal.plan_stubs.wrapped"),
122+
sources: list[
123+
Annotated[
124+
PlanSource | DeviceSource | DodalSource,
125+
Field(discriminator="kind"),
126+
]
127+
] = [
128+
PlanSource(module="dodal.plans"),
129+
PlanSource(module="dodal.plan_stubs.wrapped"),
107130
]
108131
events: WorkerEventConfig = Field(default_factory=WorkerEventConfig)
109132
metadata: MetadataConfig | None = Field(default=None)

‎src/blueapi/core/context.py‎

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,13 @@
1919

2020
from blueapi import utils
2121
from blueapi.client.numtracker import NumtrackerClient
22-
from blueapi.config import ApplicationConfig, EnvironmentConfig, SourceKind
22+
from blueapi.config import (
23+
ApplicationConfig,
24+
DeviceSource,
25+
DodalSource,
26+
EnvironmentConfig,
27+
PlanSource,
28+
)
2329
from blueapi.utils import (
2430
BlueapiPlanModelConfig,
2531
is_function_sourced_from_module,
@@ -179,14 +185,15 @@ def with_config(self, config: EnvironmentConfig) -> None:
179185
if config.metadata is not None:
180186
self.run_engine.md |= config.metadata.model_dump()
181187
for source in config.sources:
182-
mod = import_module(str(source.module))
183-
184-
if source.kind is SourceKind.PLAN_FUNCTIONS:
185-
self.with_plan_module(mod)
186-
elif source.kind is SourceKind.DEVICE_FUNCTIONS:
187-
self.with_device_module(mod)
188-
elif source.kind is SourceKind.DODAL:
189-
self.with_dodal_module(mod)
188+
mod = import_module(source.module)
189+
190+
match source:
191+
case PlanSource():
192+
self.with_plan_module(mod)
193+
case DeviceSource():
194+
self.with_device_module(mod)
195+
case DodalSource(mock=mock):
196+
self.with_dodal_module(mod, mock=mock)
190197

191198
def with_plan_module(self, module: ModuleType) -> None:
192199
"""

0 commit comments

Comments
 (0)