FilesExpand file tree

main

/

test_image_generation.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

85 lines (73 loc) · 2.97 KB

main

/

test_image_generation.py

Copy path

Top

File metadata and controls

• Code
• Blame

85 lines (73 loc) · 2.97 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

# Copyright (c) Microsoft Corporation. All rights reserved.

"""Image generation request serialization without starting a runtime."""

import json

from unittest.mock import AsyncMock, Mock

import pytest

from copilot import CopilotClient, ImageGenerationConfig, PermissionHandler, RuntimeConnection

from copilot.rpc import ProtocolImageGenerationOptions

@pytest.mark.parametrize("enabled", [None, False, True])

@pytest.mark.parametrize("watermark", [None, False, True])

def test_protocol_image_generation_preserves_positional_and_keyword_arguments(

enabled: bool | None, watermark: bool | None

):

positional = ProtocolImageGenerationOptions(enabled, watermark)

keyword = ProtocolImageGenerationOptions(enabled=enabled, watermark=watermark)

expected = {

key: value

for key, value in {"enabled": enabled, "watermark": watermark}.items()

if value is not None

}

assert positional == keyword

assert positional.to_dict() == expected

assert ProtocolImageGenerationOptions.from_dict(expected) == positional

@pytest.mark.parametrize("enabled", [None, False, True])

def test_protocol_image_generation_preserves_one_argument_constructor(enabled: bool | None):

options = ProtocolImageGenerationOptions(enabled)

assert options.enabled is enabled

assert options.watermark is None

assert options.to_dict() == ({} if enabled is None else {"enabled": enabled})

@pytest.mark.asyncio

@pytest.mark.parametrize(

"option",

[

None,

{},

{"enabled": True},

{"enabled": False},

{"watermark": True},

{"watermark": False},

{"enabled": False, "watermark": True},

{"enabled": True, "watermark": False},

],

)

async def test_image_generation_create_and_resume_wire(option: ImageGenerationConfig | None):

client = CopilotClient(connection=RuntimeConnection.for_stdio(path="unused-cli"))

client._state = "connected"

transport = Mock()

captured = {}

async def request(method, params, **kwargs):

captured[method] = json.loads(json.dumps(params))

response = {"sessionId": params["sessionId"]}

if callback := kwargs.get("on_response_inline"):

callback(response)

return response

transport.request = AsyncMock(side_effect=request)

transport.stop = AsyncMock()

client._client = transport

try:

session = await client.create_session(

on_permission_request=PermissionHandler.approve_all,

image_generation=option,

)

await client.resume_session(

session.session_id,

on_permission_request=PermissionHandler.approve_all,

image_generation=option,

)

for method in ("session.create", "session.resume"):

if option is None:

assert "imageGeneration" not in captured[method], captured[method]

else:

assert captured[method]["imageGeneration"] == option

finally:

await client.force_stop()