Skip to content

Commit

Permalink
model_test.py: Refactor test data.
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 658104296
  • Loading branch information
@schwehr
schwehr authored and Google Earth Engine Authors committed Jul 31, 2024
1 parent 5142e16 commit 82a71a3
Showing 1 changed file with 67 additions and 150 deletions.
217 changes: 67 additions & 150 deletions python/ee/tests/model_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,59 @@ def make_expression_graph(
}


def make_override_expression(key: str, pixel_type: str) -> Dict[str, Any]:
return {
'dictionaryValue': {
'values': {
key: {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': pixel_type,
'arguments': {},
}
}
},
}
}
}
}
}


def make_type_expression(
key: str, pixel_type: str, dimensions: int
) -> Dict[str, Any]:
return {
'dictionaryValue': {
'values': {
key: {
'dictionaryValue': {
'values': {
'dimensions': {'constantValue': dimensions},
'type': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': pixel_type,
'arguments': {},
}
}
},
}
},
}
}
}
}
}
}


class ModelTest(apitestcase.ApiTestCase):

def test_serialize(self):
Expand Down Expand Up @@ -88,83 +141,15 @@ def test_from_ai_platform_predictor(self):
'inputProperties': {'constantValue': input_properties},
'inputShapes': {'constantValue': input_shapes},
'inputTileSize': {'constantValue': input_tile_size},
'inputTypeOverride': {
'dictionaryValue': {
'values': {
'c': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': 'PixelType.int8',
'arguments': {},
}
}
},
}
}
}
}
},
'inputTypeOverride': make_override_expression(
'c', 'PixelType.int8'
),
'modelName': {'constantValue': model_name},
'outputBands': {
'dictionaryValue': {
'values': {
'e': {
'dictionaryValue': {
'values': {
'dimensions': {'constantValue': 10},
'type': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': (
'PixelType.int16'
),
'arguments': {},
}
}
},
}
},
}
}
}
}
}
},
'outputBands': make_type_expression('e', 'PixelType.int16', 10),
'outputMultiplier': {'constantValue': output_multiplier},
'outputProperties': {
'dictionaryValue': {
'values': {
'f': {
'dictionaryValue': {
'values': {
'dimensions': {'constantValue': 11},
'type': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': (
'PixelType.int32'
),
'arguments': {},
}
}
},
}
},
}
}
}
}
}
},
'outputProperties': make_type_expression(
'f', 'PixelType.int32', 11
),
'outputTileSize': {'constantValue': output_tile_size},
'proj': {
'functionInvocationValue': {
Expand Down Expand Up @@ -249,83 +234,15 @@ def test_from_vertex_ai(self):
'inputProperties': {'constantValue': input_properties},
'inputShapes': {'constantValue': input_shapes},
'inputTileSize': {'constantValue': input_tile_size},
'inputTypeOverride': {
'dictionaryValue': {
'values': {
'c': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': 'PixelType.int8',
'arguments': {},
}
}
},
}
}
}
}
},
'inputTypeOverride': make_override_expression(
'c', 'PixelType.int8'
),
'maxPayloadBytes': {'constantValue': max_payload_bytes},
'outputBands': {
'dictionaryValue': {
'values': {
'e': {
'dictionaryValue': {
'values': {
'dimensions': {'constantValue': 10},
'type': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': (
'PixelType.int16'
),
'arguments': {},
}
}
},
}
},
}
}
}
}
}
},
'outputBands': make_type_expression('e', 'PixelType.int16', 10),
'outputMultiplier': {'constantValue': output_multiplier},
'outputProperties': {
'dictionaryValue': {
'values': {
'f': {
'dictionaryValue': {
'values': {
'dimensions': {'constantValue': 11},
'type': {
'functionInvocationValue': {
'functionName': 'PixelType',
'arguments': {
'precision': {
'functionInvocationValue': {
'functionName': (
'PixelType.int32'
),
'arguments': {},
}
}
},
}
},
}
}
}
}
}
},
'outputProperties': make_type_expression(
'f', 'PixelType.int32', 11
),
'outputTileSize': {'constantValue': output_tile_size},
'payloadFormat': {'constantValue': payload_format},
'proj': {
Expand Down

0 comments on commit 82a71a3

Please sign in to comment.