From 7f3c7faafa1a3ef77484c4eed4f699d5fbe9db58 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@elizas.website>
Date: Thu, 6 Jun 2024 14:29:11 -0700
Subject: [PATCH] add test that migration state updates

---
 nexus/tests/integration_tests/instances.rs | 60 +++++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

diff --git a/nexus/tests/integration_tests/instances.rs b/nexus/tests/integration_tests/instances.rs
index 565e2fbafb..9ced72e2c7 100644
--- a/nexus/tests/integration_tests/instances.rs
+++ b/nexus/tests/integration_tests/instances.rs
@@ -673,6 +673,33 @@ async fn test_instance_start_creates_networking_state(
 
 #[nexus_test]
 async fn test_instance_migrate(cptestctx: &ControlPlaneTestContext) {
+    use nexus_db_model::Migration;
+    use omicron_common::api::internal::nexus::MigrationState;
+    async fn migration_fetch(
+        cptestctx: &ControlPlaneTestContext,
+        migration_id: Uuid,
+    ) -> Migration {
+        use async_bb8_diesel::AsyncRunQueryDsl;
+        use diesel::prelude::*;
+        use nexus_db_queries::db::schema::migration::dsl;
+
+        let datastore =
+            cptestctx.server.server_context().nexus.datastore().clone();
+        let db_state = dsl::migration
+            .filter(dsl::id.eq(migration_id))
+            .select(Migration::as_select())
+            .get_results_async::<Migration>(
+                &*datastore.pool_connection_for_tests().await.unwrap(),
+            )
+            .await
+            .unwrap();
+
+        info!(&cptestctx.logctx.log, "refetched migration info from db";
+                "migration" => ?db_state);
+
+        db_state.into_iter().next().unwrap()
+    }
+
     let client = &cptestctx.external_client;
     let apictx = &cptestctx.server.server_context();
     let nexus = &apictx.nexus;
@@ -729,7 +756,7 @@ async fn test_instance_migrate(cptestctx: &ControlPlaneTestContext) {
 
     let migrate_url =
         format!("/v1/instances/{}/migrate", &instance_id.to_string());
-    let _ = NexusRequest::new(
+    let instance = NexusRequest::new(
         RequestBuilder::new(client, Method::POST, &migrate_url)
             .body(Some(&params::InstanceMigrate { dst_sled_id }))
             .expect_status(Some(StatusCode::OK)),
@@ -749,6 +776,30 @@ async fn test_instance_migrate(cptestctx: &ControlPlaneTestContext) {
 
     assert_eq!(current_sled, original_sled);
 
+    // Ensure that both sled agents report that the migration is in progress.
+    let migration_id = {
+        let datastore = apictx.nexus.datastore();
+        let opctx = OpContext::for_tests(
+            cptestctx.logctx.log.new(o!()),
+            datastore.clone(),
+        );
+        let (.., authz_instance) = LookupPath::new(&opctx, &datastore)
+            .instance_id(instance.identity.id)
+            .lookup_for(nexus_db_queries::authz::Action::Read)
+            .await
+            .unwrap();
+        datastore
+            .instance_refetch(&opctx, &authz_instance)
+            .await
+            .unwrap()
+            .runtime_state
+            .migration_id
+            .expect("since we've started a migration, the instance record must have a migration id!")
+    };
+    let migration = dbg!(migration_fetch(cptestctx, migration_id).await);
+    assert_eq!(migration.target_state, MigrationState::InProgress.into());
+    assert_eq!(migration.source_state, MigrationState::InProgress.into());
+
     // Explicitly simulate the migration action on the target. Simulated
     // migrations always succeed. The state transition on the target is
     // sufficient to move the instance back into a Running state (strictly
@@ -765,6 +816,13 @@ async fn test_instance_migrate(cptestctx: &ControlPlaneTestContext) {
         .expect("migrated instance should still have a sled");
 
     assert_eq!(current_sled, dst_sled_id);
+
+    // Ensure that both sled agents report that the migration has completed.
+    instance_simulate_on_sled(cptestctx, nexus, default_sled_id, instance_id)
+        .await;
+    let migration = dbg!(migration_fetch(cptestctx, migration_id).await);
+    assert_eq!(migration.target_state, MigrationState::Completed.into());
+    assert_eq!(migration.source_state, MigrationState::Completed.into());
 }
 
 #[nexus_test]