Azure: Fix bug when updating SAS URI on existing DiskVersion #177
+362
−219
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Reviewer's GuideRefactors Azure disk version SAS URI update logic to directly modify existing VM image entries for x64 and arm64 generations, adds regression tests for multi-architecture publish flows (new vs existing DiskVersion), and removes now-unneeded VM image helper utilities and test mocks. Sequence diagram for updating SAS URI on existing DiskVersionsequenceDiagram
actor Caller
participant AzureService
participant AzureUtils
participant DiskVersion
participant VMImageDefinition as VMImage
participant VMImageSource as Source
Caller->>AzureService: publish_live(metadata, source)
AzureService->>AzureUtils: set_new_sas_disk_version(metadata, disk_version, source)
AzureUtils->>DiskVersion: check vm_images
alt disk_version has vm_images
loop for each img in disk_version.vm_images
AzureUtils->>VMImage: evaluate image_type
alt img.image_type == get_image_type_mapping(metadata.architecture, metadata.generation)
AzureUtils->>VMImage: img.source.os_disk.uri = source.os_disk.uri
else alt metadata.support_legacy and img.image_type == get_image_type_mapping(metadata.architecture, V1)
AzureUtils->>VMImage: img.source.os_disk.uri = source.os_disk.uri
else
AzureUtils-->>VMImage: leave uri unchanged
end
end
else disk_version has no vm_images
AzureUtils->>AzureUtils: create_vm_image_definitions(metadata, source)
AzureUtils->>DiskVersion: disk_version.vm_images = new_definitions
end
AzureUtils-->>AzureService: updated disk_version
AzureService-->>Caller: publish_live result
Updated class diagram for Azure VM image SAS update logicclassDiagram
class AzurePublishingMetadata {
string architecture
string generation
bool support_legacy
string image_path
}
class VMImageSource {
VMIOSDisk source_os_disk
to_json()
}
class VMIOSDisk {
string uri
}
class VMImageDefinition {
string image_type
VMImageSource source
from_json(json_data)
}
class DiskVersion {
string version_number
List~VMImageDefinition~ vm_images
}
class AzureUtils {
+set_new_sas_disk_version(metadata, disk_version, source)
+create_vm_image_definitions(metadata, source)
+get_image_type_mapping(architecture, generation)
+_all_skus_present(old_skus, disk_versions)
+seek_disk_version(disk_versions, version_number)
}
AzurePublishingMetadata --> DiskVersion : configures
AzurePublishingMetadata --> VMImageDefinition : determines image_type
VMImageSource --> VMIOSDisk : has
VMImageDefinition --> VMImageSource : has
DiskVersion --> VMImageDefinition : contains
AzureUtils ..> AzurePublishingMetadata : uses
AzureUtils ..> DiskVersion : updates
AzureUtils ..> VMImageDefinition : inspects_and_modifies
AzureUtils ..> VMImageSource : uses
%% Removed helpers (for context)
class PrepareVmImagesHelper {
-prepare_vm_images(metadata, gen1, gen2, source)
}
class VmImagesByGenerationHelper {
-vm_images_by_generation(disk_version, architecture)
}
AzureUtils ..x PrepareVmImagesHelper : removed
AzureUtils ..x VmImagesByGenerationHelper : removed
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
This commits adds the following unit-tests to AzureService: - **test_publish_live_multiple_arches_new_disk_version**: This test is passing and simulates what happens when a new DiskVersion should be created and existing ones untouched. - **test_publish_live_multiple_arches_existing_disk_version**: This test is NOT passing ATM and indicates a BUG! It simulates the behavior of replacing a single SAS URI when the DiskVersion exists. Refers to SPSTRAT-585 Signed-off-by: Jonathan Gangi <[email protected]>
This commit fixes a bug during the SAS URI update on existing DiskVersions. In the previous version it would fail to replace the required SAS URI on the proper OSDisk as it was initially developed to only deal with x86 Gen1 or Gen2 Now it properly addresses Arm64Gen2 alongside it and has a more generic algorythm to properly update the SAS in the required source. Refers to SPSTRAT-585 Signed-off-by: Jonathan Gangi <[email protected]>
No longer required after changes on `utils.set_new_sas_disk_version` Refers to SPSTRAT-585 Signed-off-by: Jonathan Gangi <[email protected]>
No longer required after changes on `utils.set_new_sas_disk_version` Refers to SPSTRAT-585 Signed-off-by: Jonathan Gangi <[email protected]>
JAVGan
changed the title
Collaborator
Author
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- In the new tests,
metadata_azure_obj.architectureis set to both"aarch64"(new disk version test) and"arm64"(existing disk version test); double-check thatget_image_type_mappingand other consumers treat these values equivalently or standardize on a single architecture string to avoid subtle mismatches. - The new
set_new_sas_disk_versionimplementation mutatesdisk_version.vm_imagesin place instead of rebuilding the list; confirm that no callers rely on the previous behavior ofvm_images_by_generation(e.g., ordering or list identity) and consider adding a brief comment explaining the in-place update for future maintainers. - The two new
test_publish_live_multiple_arches_*tests share a substantial amount of setup boilerplate (vm_versions, technical_config, product_json, HTTP mocks); extracting common helpers/fixtures for this Azure product/disk setup would make the tests easier to maintain.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In the new tests, `metadata_azure_obj.architecture` is set to both `"aarch64"` (new disk version test) and `"arm64"` (existing disk version test); double-check that `get_image_type_mapping` and other consumers treat these values equivalently or standardize on a single architecture string to avoid subtle mismatches.
- The new `set_new_sas_disk_version` implementation mutates `disk_version.vm_images` in place instead of rebuilding the list; confirm that no callers rely on the previous behavior of `vm_images_by_generation` (e.g., ordering or list identity) and consider adding a brief comment explaining the in-place update for future maintainers.
- The two new `test_publish_live_multiple_arches_*` tests share a substantial amount of setup boilerplate (vm_versions, technical_config, product_json, HTTP mocks); extracting common helpers/fixtures for this Azure product/disk setup would make the tests easier to maintain.
## Individual Comments
### Comment 1
<location> `cloudpub/ms_azure/utils.py:498-500` </location>
<code_context>
- gen2=img,
- source=source,
- )
+ for img in disk_version.vm_images:
+ if img.image_type == get_image_type_mapping(metadata.architecture, metadata.generation):
+ img.source.os_disk.uri = source.os_disk.uri
+ elif metadata.support_legacy and img.image_type == get_image_type_mapping(
+ metadata.architecture, "V1"
</code_context>
<issue_to_address>
**issue (bug_risk):** Only updating `os_disk.uri` may miss other required changes in `source`.
Previously `prepare_vm_images` rebuilt each `VMImageDefinition` from `source.to_json()`, effectively replacing the whole `source` and any derived state. Now we only update `img.source.os_disk.uri`, so other changes to `VMImageSource` won’t propagate. If `source` holds additional state (e.g. SAS params, storage account), this can desync `metadata`/`source` from `disk_version.vm_images`. Please either assign the full `source` (or all relevant fields) to `img.source`, or clearly define which subfields must stay in sync to preserve the prior behavior.
</issue_to_address>
### Comment 2
<location> `tests/ms_azure/test_service.py:1767-1699` </location>
<code_context>
+ @pytest.mark.parametrize(
</code_context>
<issue_to_address>
**suggestion (testing):** Add a test case where legacy (V1) images are updated alongside V2 for x64
Since `set_new_sas_disk_version` has special handling for `metadata.support_legacy` (updating both V2 and V1 images), please extend `test_publish_live_multiple_arches_existing_disk_version` with a parametrized case where:
- `metadata.architecture == "x64"`
- `metadata.generation == "V2"`
- `metadata.support_legacy is True`
- the disk version includes both `x64Gen2` and `x64Gen1` images
Then assert that both Gen2 and Gen1 `osDisk.uri` values are updated to the new SAS. This will directly exercise the legacy-path logic and protect against regressions when both generations are present.
Suggested implementation:
```python
@pytest.mark.parametrize(
"new_vm_image,expected_vm_images",
[
(
{
"source": {
"sourceType": "sasUri",
"osDisk": {"uri": "https://uri.test.com/newimg-2.0-20260101.0.aarch64.vhd"},
"dataDisks": [],
},
"imageType": "arm64Gen2",
```
To fully implement the requested test coverage, please adjust the same parametrized block and the test body as follows (you will see the exact surrounding code in your local file):
1. **Extend the parametrization with a legacy-support x64 case**
Inside the list passed to `@pytest.mark.parametrize("new_vm_image,expected_vm_images", [...])`, append a new tuple that represents:
- A new x64Gen2 image using a new SAS URI.
- An expected disk version that contains both `x64Gen2` and `x64Gen1` images, *both* using that new SAS URI in their `osDisk.uri`.
For example, at the end of the existing list of `(new_vm_image, expected_vm_images)` tuples, add something like:
```python
# legacy (V1) + V2 images for x64 should both get the new SAS URI
(
{
"source": {
"sourceType": "sasUri",
"osDisk": {
"uri": "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd"
},
"dataDisks": [],
},
"imageType": "x64Gen2",
},
[
{
"source": {
"sourceType": "sasUri",
"osDisk": {
"uri": "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd"
},
"dataDisks": [],
},
"imageType": "x64Gen2",
},
{
"source": {
"sourceType": "sasUri",
"osDisk": {
"uri": "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd"
},
"dataDisks": [],
},
"imageType": "x64Gen1",
},
],
),
```
Adjust the URI pattern to match the conventions used by your other test cases (e.g. `x86_64` vs `x64`, date stamp format, etc.).
2. **Ensure metadata drives the legacy behaviour**
In `test_publish_live_multiple_arches_existing_disk_version`:
- Make sure the test either:
- Accepts a `metadata` parameter from another `@pytest.mark.parametrize` or fixture, **or**
- Constructs the metadata inline.
- For the new case, ensure the metadata passed into the code under test has:
```python
metadata.architecture == "x64"
metadata.generation == "V2"
metadata.support_legacy is True
```
How to do this will depend on your existing test pattern; commonly you'll have something like a `metadata` object or dict passed into the service being tested. For this new parameter set, you may need to:
- Add an extra argument to the existing `@pytest.mark.parametrize` (e.g. `"new_vm_image, expected_vm_images, metadata"`), and update existing cases, **or**
- Add conditional logic inside the test that sets `metadata.support_legacy = True` when `new_vm_image["imageType"] == "x64Gen2"` and/or when an additional parameter (like `support_legacy`) is True.
3. **Assert both Gen2 and Gen1 URIs are updated**
After calling the function that eventually invokes `set_new_sas_disk_version`, you likely have an assertion comparing the resulting disk version VM images to `expected_vm_images`.
If you assert on the full VM images structure already (e.g. `assert vm_image_list == expected_vm_images`), your new parameter set will automatically verify that both `x64Gen2` and `x64Gen1` `osDisk.uri` values match the new SAS URI.
If instead you assert more granularly, add explicit checks such as:
```python
# Pseudocode – adapt to your actual result structure/variable names
x64_gen2 = next(i for i in resulting_vm_images if i["imageType"] == "x64Gen2")
x64_gen1 = next(i for i in resulting_vm_images if i["imageType"] == "x64Gen1")
assert x64_gen2["source"]["osDisk"]["uri"] == "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd"
assert x64_gen1["source"]["osDisk"]["uri"] == "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd"
```
4. **Keep the test name and purpose intact**
No need to rename `test_publish_live_multiple_arches_existing_disk_version`; the added case will keep its current semantics but additionally validate the legacy-support path where both Gen2 and Gen1 images are present and should be updated together.
Because I only see a small slice of the parametrization and no body of the test, you'll need to integrate the snippet above into the full list and adjust variable names/types (e.g. `x64` vs `x86_64`, how `metadata` is constructed) to match your existing test code.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+498
to
+500
| for img in disk_version.vm_images: | ||
| if img.image_type == get_image_type_mapping(metadata.architecture, metadata.generation): | ||
| img.source.os_disk.uri = source.os_disk.uri |
There was a problem hiding this comment.
issue (bug_risk): Only updating os_disk.uri may miss other required changes in source.
Previously prepare_vm_images rebuilt each VMImageDefinition from source.to_json(), effectively replacing the whole source and any derived state. Now we only update img.source.os_disk.uri, so other changes to VMImageSource won’t propagate. If source holds additional state (e.g. SAS params, storage account), this can desync metadata/source from disk_version.vm_images. Please either assign the full source (or all relevant fields) to img.source, or clearly define which subfields must stay in sync to preserve the prior behavior.
| }, | ||
| "imageType": "arm64Gen2", | ||
| }, | ||
| ], |
There was a problem hiding this comment.
suggestion (testing): Add a test case where legacy (V1) images are updated alongside V2 for x64
Since set_new_sas_disk_version has special handling for metadata.support_legacy (updating both V2 and V1 images), please extend test_publish_live_multiple_arches_existing_disk_version with a parametrized case where:
metadata.architecture == "x64"metadata.generation == "V2"metadata.support_legacy is True- the disk version includes both
x64Gen2andx64Gen1images
Then assert that both Gen2 and Gen1 osDisk.uri values are updated to the new SAS. This will directly exercise the legacy-path logic and protect against regressions when both generations are present.
Suggested implementation:
@pytest.mark.parametrize(
"new_vm_image,expected_vm_images",
[
(
{
"source": {
"sourceType": "sasUri",
"osDisk": {"uri": "https://uri.test.com/newimg-2.0-20260101.0.aarch64.vhd"},
"dataDisks": [],
},
"imageType": "arm64Gen2",To fully implement the requested test coverage, please adjust the same parametrized block and the test body as follows (you will see the exact surrounding code in your local file):
-
Extend the parametrization with a legacy-support x64 case
Inside the list passed to
@pytest.mark.parametrize("new_vm_image,expected_vm_images", [...]), append a new tuple that represents:- A new x64Gen2 image using a new SAS URI.
- An expected disk version that contains both
x64Gen2andx64Gen1images, both using that new SAS URI in theirosDisk.uri.
For example, at the end of the existing list of
(new_vm_image, expected_vm_images)tuples, add something like:# legacy (V1) + V2 images for x64 should both get the new SAS URI ( { "source": { "sourceType": "sasUri", "osDisk": { "uri": "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd" }, "dataDisks": [], }, "imageType": "x64Gen2", }, [ { "source": { "sourceType": "sasUri", "osDisk": { "uri": "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd" }, "dataDisks": [], }, "imageType": "x64Gen2", }, { "source": { "sourceType": "sasUri", "osDisk": { "uri": "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd" }, "dataDisks": [], }, "imageType": "x64Gen1", }, ], ),
Adjust the URI pattern to match the conventions used by your other test cases (e.g.
x86_64vsx64, date stamp format, etc.). -
Ensure metadata drives the legacy behaviour
In
test_publish_live_multiple_arches_existing_disk_version:- Make sure the test either:
- Accepts a
metadataparameter from another@pytest.mark.parametrizeor fixture, or - Constructs the metadata inline.
- Accepts a
- For the new case, ensure the metadata passed into the code under test has:
metadata.architecture == "x64" metadata.generation == "V2" metadata.support_legacy is True
How to do this will depend on your existing test pattern; commonly you'll have something like a
metadataobject or dict passed into the service being tested. For this new parameter set, you may need to:- Add an extra argument to the existing
@pytest.mark.parametrize(e.g."new_vm_image, expected_vm_images, metadata"), and update existing cases, or - Add conditional logic inside the test that sets
metadata.support_legacy = Truewhennew_vm_image["imageType"] == "x64Gen2"and/or when an additional parameter (likesupport_legacy) is True.
- Make sure the test either:
-
Assert both Gen2 and Gen1 URIs are updated
After calling the function that eventually invokes
set_new_sas_disk_version, you likely have an assertion comparing the resulting disk version VM images toexpected_vm_images.If you assert on the full VM images structure already (e.g.
assert vm_image_list == expected_vm_images), your new parameter set will automatically verify that bothx64Gen2andx64Gen1osDisk.urivalues match the new SAS URI.If instead you assert more granularly, add explicit checks such as:
# Pseudocode – adapt to your actual result structure/variable names x64_gen2 = next(i for i in resulting_vm_images if i["imageType"] == "x64Gen2") x64_gen1 = next(i for i in resulting_vm_images if i["imageType"] == "x64Gen1") assert x64_gen2["source"]["osDisk"]["uri"] == "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd" assert x64_gen1["source"]["osDisk"]["uri"] == "https://uri.test.com/newimg-2.0-20260101.0.x64.vhd"
-
Keep the test name and purpose intact
No need to rename
test_publish_live_multiple_arches_existing_disk_version; the added case will keep its current semantics but additionally validate the legacy-support path where both Gen2 and Gen1 images are present and should be updated together.
Because I only see a small slice of the parametrization and no body of the test, you'll need to integrate the snippet above into the full list and adjust variable names/types (e.g. x64 vs x86_64, how metadata is constructed) to match your existing test code.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Fix bug when updating SAS URI on existing DiskVersion
This PR contains the following changes:
Add unit tests to reproduce bug on updating DV
This commit adds the following unit-tests to AzureService in order to reproduce the issue:
test_publish_live_multiple_arches_new_disk_version: This test is passing and simulates what happens when a new DiskVersion should be created and existing ones untouched.
test_publish_live_multiple_arches_existing_disk_version: This test was not passing with the old code and indicates a BUG! It simulates the behavior of replacing a single SAS URI when the DiskVersion exists.
Properly address SAS during DV Update
This commit fixes a bug during the SAS URI update on existing DiskVersions.
In the previous version it would fail to replace the required SAS URI on the proper OSDisk as it was initially developed to only deal with x86 Gen1 or Gen2
Now it properly addresses Arm64Gen2 alongside it and has a more generic algorithm to properly update the SAS in the required source.
Deimplement "utils.prepare_vm_images"
Old code no longer required after the fix - simplifies the module.
Deimplement "utils.vm_images_by_generation"
Old code no longer required after the fix - simplifies the module.
JIRA
Refers to SPSTRAT-585
Summary by Sourcery
Fix Azure disk version SAS URI updates across architectures and simplify related utilities and tests.
Bug Fixes:
Enhancements:
Tests: