Merge pull request #21 from fermyon/custom-services

Custom services

Author: rylev

crates/test-environment/src/services.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     path::{Path, PathBuf},
 };
 
@@ -11,6 +11,8 @@ use anyhow::{bail, Context};
 use docker::DockerService;
 use python::PythonService;
 
+pub use docker::DockerImage;
+
 /// All the services that are running for a test.
 #[derive(Default)]
 pub struct Services {
@@ -19,25 +21,21 @@ pub struct Services {
 
 impl Services {
     /// Start all the required services given a path to service definitions
-    pub fn start(config: &ServicesConfig, working_dir: &Path) -> anyhow::Result<Self> {
+    pub fn start(config: ServicesConfig, working_dir: &Path) -> anyhow::Result<Self> {
+        let lock_dir = working_dir.join(".service-locks");
+        std::fs::create_dir(&lock_dir).context("could not create service lock dir")?;
         let mut services = Vec::new();
-        for required_service in &config.services {
-            let service_definition_extension =
-                config.definitions.get(required_service).map(|e| e.as_str());
-            let mut service: Box<dyn Service> = match service_definition_extension {
-                Some("py") => Box::new(PythonService::start(
-                    required_service,
-                    &config.definitions_path,
+        for service_def in config.service_definitions {
+            let mut service: Box<dyn Service> = match service_def.kind {
+                ServiceKind::Python { script } => Box::new(PythonService::start(
+                    &service_def.name,
+                    &script,
                     working_dir,
+                    &lock_dir,
                 )?),
-                Some("Dockerfile") => Box::new(DockerService::start(
-                    required_service,
-                    &config.definitions_path,
-                )?),
-                Some(extension) => {
-                    bail!("service definitions with the '{extension}' extension are not supported")
+                ServiceKind::Docker { image } => {
+                    Box::new(DockerService::start(&service_def.name, image, &lock_dir)?)
                 }
-                None => bail!("no service definition found for '{required_service}'"),
             };
             service.ready()?;
             services.push(service);
@@ -88,37 +86,46 @@ impl<'a> IntoIterator for &'a Services {
 }
 
 pub struct ServicesConfig {
-    services: Vec<String>,
-    definitions_path: PathBuf,
-    definitions: HashMap<String, String>,
+    /// Definitions of all services to be used.
+    service_definitions: Vec<ServiceDefinition>,
 }
 
 impl ServicesConfig {
-    /// Create a new services config a list of services to start.
+    /// Create a new services config with a list of built-in services to start.
     ///
-    /// The services are expected to have a definition file in the `services` directory with the same name as the service.
-    pub fn new(services: Vec<String>) -> anyhow::Result<Self> {
-        let definitions = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("services");
-        let service_definitions = service_definitions(&definitions)?;
+    /// The built-in services are expected to have a definition file in the `services` directory with the same name as the service.
+    pub fn new<'a>(builtins: impl Into<Vec<&'a str>>) -> anyhow::Result<Self> {
+        let definitions_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("services");
+        let service_definitions = get_builtin_service_definitions(
+            builtins.into().into_iter().collect(),
+            &definitions_path,
+        )?;
         Ok(Self {
-            services,
-            definitions_path: definitions,
-            definitions: service_definitions,
+            service_definitions,
         })
     }
 
+    pub fn add_service(&mut self, service: ServiceDefinition) {
+        self.service_definitions.push(service);
+    }
+
     /// Configure no services
     pub fn none() -> Self {
         Self {
-            services: Vec::new(),
-            definitions_path: PathBuf::new(),
-            definitions: HashMap::new(),
+            service_definitions: Vec::new(),
         }
     }
 }
 
 /// Get all of the service definitions returning a HashMap of the service name to the service definition file extension.
-fn service_definitions(service_definitions_path: &Path) -> anyhow::Result<HashMap<String, String>> {
+fn get_builtin_service_definitions(
+    mut builtins: HashSet<&str>,
+    service_definitions_path: &Path,
+) -> anyhow::Result<Vec<ServiceDefinition>> {
+    if builtins.is_empty() {
+        return Ok(Vec::new());
+    }
+
     std::fs::read_dir(service_definitions_path)
         .with_context(|| {
             format!(
@@ -139,10 +146,43 @@ fn service_definitions(service_definitions_path: &Path) -> anyhow::Result<HashMa
                 .context("service definition did not have an extension")?;
             Ok((file_name.to_owned(), file_extension.to_owned()))
         })
-        .filter(|r| !matches!( r , Ok((_, extension)) if extension == "lock"))
+        .filter(|r| !matches!(r, Ok((_, extension)) if extension == "lock"))
+        .filter(move |r| match r {
+            Ok((service, _)) => builtins.remove(service.as_str()),
+            _ => false,
+        })
+        .map(|r| {
+            let (name, extension) = r?;
+            Ok(ServiceDefinition {
+                name: name.clone(),
+                kind: match extension.as_str() {
+                    "py" => ServiceKind::Python {
+                        script: service_definitions_path.join(format!("{}.py", name)),
+                    },
+                    "Dockerfile" => ServiceKind::Docker {
+                        image: docker::DockerImage::FromDockerfile(
+                            service_definitions_path.join(format!("{}.Dockerfile", name)),
+                        ),
+                    },
+                    _ => bail!("unsupported service definition extension '{}'", extension),
+                },
+            })
+        })
         .collect()
 }
 
+/// A service definition.
+pub struct ServiceDefinition {
+    pub name: String,
+    pub kind: ServiceKind,
+}
+
+/// The kind of service.
+pub enum ServiceKind {
+    Python { script: PathBuf },
+    Docker { image: DockerImage },
+}
+
 /// An external service a test may depend on.
 pub trait Service {
     /// The name of the service.

crates/test-environment/src/services/docker.rs

@@ -3,7 +3,7 @@ use anyhow::{bail, Context as _};
 use std::{
     cell::OnceCell,
     collections::HashMap,
-    path::Path,
+    path::{Path, PathBuf},
     process::{Command, Stdio},
 };
 
@@ -19,17 +19,22 @@ pub struct DockerService {
 
 impl DockerService {
     /// Start a docker container as a service
-    pub fn start(name: &str, service_definitions_path: &Path) -> anyhow::Result<Self> {
+    pub fn start(name: &str, image: DockerImage, lock_dir: &Path) -> anyhow::Result<Self> {
+        let lock_path = lock_dir.join(format!("{name}.lock"));
         // TODO: ensure that `docker` is installed and available
-        let docker_file_path = service_definitions_path.join(format!("{name}.Dockerfile"));
-        let image_name = format!("test-environment/services/{name}");
         let mut lock =
-            fslock::LockFile::open(&service_definitions_path.join(format!("{name}.lock")))
-                .context("failed to open service file lock")?;
+            fslock::LockFile::open(&lock_path).context("failed to open service file lock")?;
         lock.lock().context("failed to obtain service file lock")?;
 
+        let image_name = match image {
+            DockerImage::FromDockerfile(dockerfile_path) => {
+                let image_name = format!("test-environment/services/{name}");
+                build_image(&dockerfile_path, &image_name)?;
+                image_name
+            }
+            DockerImage::FromRegistry(image_name) => image_name,
+        };
         stop_containers(&get_running_containers(&image_name)?)?;
-        build_image(&docker_file_path, &image_name)?;
         let container = run_container(&image_name)?;
 
         Ok(Self {
@@ -78,6 +83,11 @@ impl Container {
     }
 }
 
+pub enum DockerImage {
+    FromDockerfile(PathBuf),
+    FromRegistry(String),
+}
+
 impl Drop for Container {
     fn drop(&mut self) {
         let _ = stop_containers(&[std::mem::take(&mut self.id)]);
@@ -141,13 +151,13 @@ impl Service for DockerService {
     }
 }
 
-fn build_image(docker_file_path: &Path, image_name: &String) -> anyhow::Result<()> {
+fn build_image(dockerfile_path: &Path, image_name: &String) -> anyhow::Result<()> {
     let temp_dir = temp_dir::TempDir::new()
         .context("failed to produce a temporary directory to run docker in")?;
     let output = Command::new("docker")
         .arg("build")
         .arg("-f")
-        .arg(docker_file_path)
+        .arg(dockerfile_path)
         .arg("-t")
         .arg(image_name)
         .arg(temp_dir.path())

crates/test-environment/src/services/python.rs

@@ -21,21 +21,17 @@ pub struct PythonService {
 impl PythonService {
     pub fn start(
         name: &str,
-        service_definitions_path: &Path,
+        script_path: &Path,
         working_dir: &Path,
+        lock_dir: &Path,
     ) -> anyhow::Result<Self> {
+        let lock_path = lock_dir.join(format!("{name}.lock"));
         let mut lock =
-            fslock::LockFile::open(&service_definitions_path.join(format!("{name}.lock")))
-                .context("failed to open service file lock")?;
+            fslock::LockFile::open(&lock_path).context("failed to open service file lock")?;
         lock.lock().context("failed to obtain service file lock")?;
         let mut child = python()
             .current_dir(working_dir)
-            .arg(
-                service_definitions_path
-                    .join(format!("{name}.py"))
-                    .display()
-                    .to_string(),
-            )
+            .arg(script_path.display().to_string())
             .stdout(Stdio::piped())
             .spawn()
             .context("service failed to spawn")?;

crates/test-environment/src/test_environment.rs

@@ -25,7 +25,7 @@ impl<R: Runtime> TestEnvironment<R> {
         config: TestEnvironmentConfig<R>,
         init_env: impl FnOnce(&mut Self) -> anyhow::Result<()> + 'static,
     ) -> anyhow::Result<Self> {
-        let mut env = Self::boot(&config.services_config)?;
+        let mut env = Self::boot(config.services_config)?;
         init_env(&mut env)?;
         let runtime = (config.create_runtime)(&mut env)?;
         env.start_runtime(runtime)
@@ -47,7 +47,7 @@ impl<R> TestEnvironment<R> {
     /// Spin up a test environment without a runtime
     ///
     /// `services` specifies the services to run.
-    pub fn boot(services: &ServicesConfig) -> anyhow::Result<Self> {
+    pub fn boot(services: ServicesConfig) -> anyhow::Result<Self> {
         let temp = temp_dir::TempDir::new()
             .context("failed to produce a temporary directory to run the test in")?;
         let mut services =

src/main.rs

@@ -113,7 +113,7 @@ fn package_into(dir_path: impl AsRef<Path>) -> anyhow::Result<()> {
             continue;
         }
         let test_path = test.path();
-        println!("Processing test {:?}...", test_path);
+        println!("Processing {test_path:?}...");
 
         let test_name = test_path
             .file_name()
@@ -180,7 +180,6 @@ fn substitute_source(
                     .with_context(|| format!("'{template_value}' is not a known component"))?;
                 let component_file = "component.wasm";
                 std::fs::copy(path, test_archive.join(component_file))?;
-                println!("Substituting {template} with {component_file}...");
                 manifest.replace_range(full.range(), component_file);
                 // Restart the search after a substitution
                 continue 'outer;