NOISSUE - Prepare cocos for v0.8.0 (#512)

* Refactor mock interfaces to use 'any' instead of 'interface{}' for improved type safety and readability across multiple files in the manager and pkg directories.

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update Go version to 1.25.x in CI workflows and remove obsolete Go package files

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add mock implementations for various components in the attestation and SDK packages

- Created mock for MeasurementProvider in pkg/attestation/cmdconfig/mocks/mocks_test.go
- Created mock for Provider in pkg/attestation/mocks/mocks_test.go
- Created mock for Client in pkg/clients/grpc/mocks/mocks_test.go
- Created mock for SDK in pkg/sdk/mocks/mocks_test.go

These mocks are generated using mockery and are intended for unit testing purposes.

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Remove autogenerated mock files and update mock usage in tests

- Deleted mocks for gRPC clients in pkg/clients/grpc/mocks/mocks_test.go and pkg/sdk/mocks/mocks_test.go.
- Updated test files in pkg/progressbar/progress_test.go to use the new mock structure without type parameters for gRPC client interfaces.
- Refactored mock generation in pkg/sdk/mocks/sdk.go to streamline the mock creation process and ensure consistency across mock methods.

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update protobuf generated files for events and manager

- Bump protoc-gen-go version from v1.36.5 to v1.36.8 in events.pb.go and manager.pb.go.
- Refactor raw descriptor definitions in events.pb.go and manager.pb.go to use string concatenation for better readability and maintainability.
- Ensure compatibility with the latest protobuf specifications and improve code generation consistency.

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update test commands to use GOTOOLCHAIN for consistent Go version handling

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Fix GOTOOLCHAIN usage in test command for consistency

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

agent/agent.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: agent/agent.proto
 
@@ -570,76 +570,44 @@ func (x *AttestationTokenResponse) GetFile() []byte {
 
 var File_agent_agent_proto protoreflect.FileDescriptor
 
-var file_agent_agent_proto_rawDesc = string([]byte{
-	0x0a, 0x11, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x70, 0x72,
-	0x6f, 0x74, 0x6f, 0x12, 0x05, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x22, 0x4f, 0x0a, 0x0b, 0x41, 0x6c,
-	0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1c, 0x0a, 0x09, 0x61, 0x6c, 0x67,
-	0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x61, 0x6c,
-	0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x22, 0x0a, 0x0c, 0x72, 0x65, 0x71, 0x75, 0x69,
-	0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0c, 0x72,
-	0x65, 0x71, 0x75, 0x69, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x22, 0x0e, 0x0a, 0x0c, 0x41,
-	0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x43, 0x0a, 0x0b, 0x44,
-	0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x61,
-	0x74, 0x61, 0x73, 0x65, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x61, 0x74,
-	0x61, 0x73, 0x65, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
-	0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
-	0x22, 0x0e, 0x0a, 0x0c, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x0f, 0x0a, 0x0d, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
-	0x74, 0x22, 0x24, 0x0a, 0x0e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f,
-	0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x62, 0x0a, 0x12, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1a, 0x0a,
-	0x08, 0x74, 0x65, 0x65, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52,
-	0x08, 0x74, 0x65, 0x65, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x12, 0x1c, 0x0a, 0x09, 0x76, 0x74, 0x70,
-	0x6d, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x76, 0x74,
-	0x70, 0x6d, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18,
-	0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x22, 0x29, 0x0a, 0x13, 0x41,
-	0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
-	0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c,
-	0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x18, 0x0a, 0x16, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61,
-	0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
-	0x22, 0x43, 0x0a, 0x17, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65,
-	0x6e, 0x74, 0x73, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66,
-	0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x12,
-	0x14, 0x0a, 0x05, 0x70, 0x63, 0x72, 0x31, 0x30, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x05,
-	0x70, 0x63, 0x72, 0x31, 0x30, 0x22, 0x4d, 0x0a, 0x17, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
-	0x12, 0x1e, 0x0a, 0x0a, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x0c, 0x52, 0x0a, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x4e, 0x6f, 0x6e, 0x63, 0x65,
-	0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04,
-	0x74, 0x79, 0x70, 0x65, 0x22, 0x2e, 0x0a, 0x18, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74,
-	0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04,
-	0x66, 0x69, 0x6c, 0x65, 0x32, 0xaf, 0x03, 0x0a, 0x0c, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x65,
-	0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x04, 0x41, 0x6c, 0x67, 0x6f, 0x12, 0x12, 0x2e,
-	0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
-	0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65,
-	0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28, 0x01, 0x12, 0x33, 0x0a, 0x04, 0x44, 0x61,
-	0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52,
-	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x44,
-	0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28, 0x01, 0x12,
-	0x39, 0x0a, 0x06, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x14, 0x2e, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a,
-	0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65,
-	0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x48, 0x0a, 0x0b, 0x41, 0x74,
-	0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74,
-	0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x00, 0x30, 0x01, 0x12, 0x54, 0x0a, 0x0f, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75,
-	0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x1d, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e,
-	0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52,
-	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x49,
-	0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65,
-	0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x5a, 0x0a, 0x15, 0x41, 0x7a,
-	0x75, 0x72, 0x65, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f,
-	0x6b, 0x65, 0x6e, 0x12, 0x1e, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65,
-	0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x52, 0x65, 0x71, 0x75,
-	0x65, 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65,
-	0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x52, 0x65, 0x73, 0x70,
-	0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-})
+const file_agent_agent_proto_rawDesc = "" +
+	"\n" +
+	"\x11agent/agent.proto\x12\x05agent\"O\n" +
+	"\vAlgoRequest\x12\x1c\n" +
+	"\talgorithm\x18\x01 \x01(\fR\talgorithm\x12\"\n" +
+	"\frequirements\x18\x02 \x01(\fR\frequirements\"\x0e\n" +
+	"\fAlgoResponse\"C\n" +
+	"\vDataRequest\x12\x18\n" +
+	"\adataset\x18\x01 \x01(\fR\adataset\x12\x1a\n" +
+	"\bfilename\x18\x02 \x01(\tR\bfilename\"\x0e\n" +
+	"\fDataResponse\"\x0f\n" +
+	"\rResultRequest\"$\n" +
+	"\x0eResultResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\"b\n" +
+	"\x12AttestationRequest\x12\x1a\n" +
+	"\bteeNonce\x18\x01 \x01(\fR\bteeNonce\x12\x1c\n" +
+	"\tvtpmNonce\x18\x02 \x01(\fR\tvtpmNonce\x12\x12\n" +
+	"\x04type\x18\x03 \x01(\x05R\x04type\")\n" +
+	"\x13AttestationResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\"\x18\n" +
+	"\x16IMAMeasurementsRequest\"C\n" +
+	"\x17IMAMeasurementsResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\x12\x14\n" +
+	"\x05pcr10\x18\x02 \x01(\fR\x05pcr10\"M\n" +
+	"\x17AttestationTokenRequest\x12\x1e\n" +
+	"\n" +
+	"tokenNonce\x18\x01 \x01(\fR\n" +
+	"tokenNonce\x12\x12\n" +
+	"\x04type\x18\x03 \x01(\x05R\x04type\".\n" +
+	"\x18AttestationTokenResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file2\xaf\x03\n" +
+	"\fAgentService\x123\n" +
+	"\x04Algo\x12\x12.agent.AlgoRequest\x1a\x13.agent.AlgoResponse\"\x00(\x01\x123\n" +
+	"\x04Data\x12\x12.agent.DataRequest\x1a\x13.agent.DataResponse\"\x00(\x01\x129\n" +
+	"\x06Result\x12\x14.agent.ResultRequest\x1a\x15.agent.ResultResponse\"\x000\x01\x12H\n" +
+	"\vAttestation\x12\x19.agent.AttestationRequest\x1a\x1a.agent.AttestationResponse\"\x000\x01\x12T\n" +
+	"\x0fIMAMeasurements\x12\x1d.agent.IMAMeasurementsRequest\x1a\x1e.agent.IMAMeasurementsResponse\"\x000\x01\x12Z\n" +
+	"\x15AzureAttestationToken\x12\x1e.agent.AttestationTokenRequest\x1a\x1f.agent.AttestationTokenResponse\"\x00B\tZ\a./agentb\x06proto3"
 
 var (
 	file_agent_agent_proto_rawDescOnce sync.Once

agent/algorithm/mocks/algorithm.go

@@ -1,11 +1,29 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
-import mock "github.com/stretchr/testify/mock"
+import (
+	mock "github.com/stretchr/testify/mock"
+)
+
+// NewAlgorithm creates a new instance of Algorithm. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAlgorithm(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Algorithm {
+	mock := &Algorithm{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
 
 // Algorithm is an autogenerated mock type for the Algorithm type
 type Algorithm struct {
@@ -20,21 +38,20 @@ func (_m *Algorithm) EXPECT() *Algorithm_Expecter {
 	return &Algorithm_Expecter{mock: &_m.Mock}
 }
 
-// Run provides a mock function with no fields
-func (_m *Algorithm) Run() error {
-	ret := _m.Called()
+// Run provides a mock function for the type Algorithm
+func (_mock *Algorithm) Run() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Run")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -55,8 +72,8 @@ func (_c *Algorithm_Run_Call) Run(run func()) *Algorithm_Run_Call {
 	return _c
 }
 
-func (_c *Algorithm_Run_Call) Return(_a0 error) *Algorithm_Run_Call {
-	_c.Call.Return(_a0)
+func (_c *Algorithm_Run_Call) Return(err error) *Algorithm_Run_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -65,21 +82,20 @@ func (_c *Algorithm_Run_Call) RunAndReturn(run func() error) *Algorithm_Run_Call
 	return _c
 }
 
-// Stop provides a mock function with no fields
-func (_m *Algorithm) Stop() error {
-	ret := _m.Called()
+// Stop provides a mock function for the type Algorithm
+func (_mock *Algorithm) Stop() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Stop")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -100,8 +116,8 @@ func (_c *Algorithm_Stop_Call) Run(run func()) *Algorithm_Stop_Call {
 	return _c
 }
 
-func (_c *Algorithm_Stop_Call) Return(_a0 error) *Algorithm_Stop_Call {
-	_c.Call.Return(_a0)
+func (_c *Algorithm_Stop_Call) Return(err error) *Algorithm_Stop_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -109,17 +125,3 @@ func (_c *Algorithm_Stop_Call) RunAndReturn(run func() error) *Algorithm_Stop_Ca
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewAlgorithm creates a new instance of Algorithm. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewAlgorithm(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Algorithm {
-	mock := &Algorithm{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/api/grpc/endpoint.go

@@ -11,7 +11,7 @@ import (
 )
 
 func algoEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(algoReq)
 
 		if err := req.validate(); err != nil {
@@ -30,7 +30,7 @@ func algoEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func dataEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(dataReq)
 
 		if err := req.validate(); err != nil {
@@ -49,7 +49,7 @@ func dataEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func resultEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(resultReq)
 
 		if err := req.validate(); err != nil {
@@ -65,7 +65,7 @@ func resultEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(attestationReq)
 
 		if err := req.validate(); err != nil {
@@ -81,7 +81,7 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(imaMeasurementsReq)
 
 		if err := req.validate(); err != nil {
@@ -97,7 +97,7 @@ func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func azureAttestationTokenEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(azureAttestationTokenReq)
 		if err := req.validate(); err != nil {
 			return fetchAttestationTokenRes{}, err

agent/api/grpc/interceptors.go

@@ -31,7 +31,7 @@ func NewAuthInterceptor(authSvc auth.Authenticator) (grpc.UnaryServerInterceptor
 }
 
 func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
-	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		switch info.FullMethod {
 		case agent.AgentService_Algo_FullMethodName:
 			if _, err := s.auth.AuthenticateUser(stream.Context(), auth.AlgorithmProviderRole); err != nil {
@@ -59,7 +59,7 @@ func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
 }
 
 func (s *authInterceptor) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 		switch info.FullMethod {
 		case agent.AgentService_Result_FullMethodName:
 			ctx, err := s.auth.AuthenticateUser(ctx, auth.ConsumerRole)

agent/api/grpc/interceptors_test.go

@@ -58,7 +58,7 @@ func TestAuthUnaryInterceptor(t *testing.T) {
 			}
 			unaryInt, _ := NewAuthInterceptor(authmock)
 
-			_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req interface{}) (interface{}, error) {
+			_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req any) (any, error) {
 				return nil, nil
 			})
 
@@ -129,7 +129,7 @@ func TestAuthStreamInterceptor(t *testing.T) {
 			}
 			_, streamInt := NewAuthInterceptor(authmock)
 
-			err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv interface{}, stream grpc.ServerStream) error {
+			err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv any, stream grpc.ServerStream) error {
 				return nil
 			})
 

agent/api/grpc/server.go

@@ -96,7 +96,7 @@ func NewServer(svc agent.Service) agent.AgentServiceServer {
 	}
 }
 
-func decodeAlgoRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeAlgoRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.AlgoRequest)
 	return algoReq{
 		Algorithm:    req.Algorithm,
@@ -104,11 +104,11 @@ func decodeAlgoRequest(_ context.Context, grpcReq interface{}) (interface{}, err
 	}, nil
 }
 
-func encodeAlgoResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeAlgoResponse(_ context.Context, response any) (any, error) {
 	return &agent.AlgoResponse{}, nil
 }
 
-func decodeDataRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeDataRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.DataRequest)
 	return dataReq{
 		Dataset:  req.Dataset,
@@ -116,22 +116,22 @@ func decodeDataRequest(_ context.Context, grpcReq interface{}) (interface{}, err
 	}, nil
 }
 
-func encodeDataResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeDataResponse(_ context.Context, response any) (any, error) {
 	return &agent.DataResponse{}, nil
 }
 
-func decodeResultRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeResultRequest(_ context.Context, grpcReq any) (any, error) {
 	return resultReq{}, nil
 }
 
-func encodeResultResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeResultResponse(_ context.Context, response any) (any, error) {
 	res := response.(resultRes)
 	return &agent.ResultResponse{
 		File: res.File,
 	}, nil
 }
 
-func validateNonce(nonce []byte, maxLen int, target interface{}) error {
+func validateNonce(nonce []byte, maxLen int, target any) error {
 	if len(nonce) > maxLen {
 		switch maxLen {
 		case quoteprovider.Nonce:
@@ -154,7 +154,7 @@ func validateNonce(nonce []byte, maxLen int, target interface{}) error {
 	return nil
 }
 
-func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeAttestationRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.AttestationRequest)
 	var reportData [quoteprovider.Nonce]byte
 	var nonce [vtpm.Nonce]byte
@@ -174,14 +174,14 @@ func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface
 	}, nil
 }
 
-func encodeAttestationResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeAttestationResponse(_ context.Context, response any) (any, error) {
 	res := response.(attestationRes)
 	return &agent.AttestationResponse{
 		File: res.File,
 	}, nil
 }
 
-func decodeAttestationTokenRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeAttestationTokenRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.AttestationTokenRequest)
 	var nonce [vtpm.Nonce]byte
 
@@ -193,18 +193,18 @@ func decodeAttestationTokenRequest(_ context.Context, grpcReq interface{}) (inte
 	}, nil
 }
 
-func encodeAttestationTokenResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeAttestationTokenResponse(_ context.Context, response any) (any, error) {
 	res := response.(fetchAttestationTokenRes)
 	return &agent.AttestationTokenResponse{
 		File: res.File,
 	}, nil
 }
 
-func decodeIMAMeasurementsRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeIMAMeasurementsRequest(_ context.Context, grpcReq any) (any, error) {
 	return imaMeasurementsReq{}, nil
 }
 
-func encodeIMAMeasurementsResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeIMAMeasurementsResponse(_ context.Context, response any) (any, error) {
 	res := response.(imaMeasurementsRes)
 	return &agent.IMAMeasurementsResponse{
 		File:  res.File,
@@ -215,10 +215,10 @@ func encodeIMAMeasurementsResponse(_ context.Context, response interface{}) (int
 func (s *grpcServer) streamingHandler(
 	ctx context.Context,
 	handlerName string,
-	req interface{},
-	stream interface{},
+	req any,
+	stream any,
 	sendFn func([]byte) error,
-	getFileData func(interface{}) []byte,
+	getFileData func(any) []byte,
 ) error {
 	handler, ok := s.handlers[handlerName]
 	if !ok {
@@ -350,7 +350,7 @@ func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_
 		func(data []byte) error {
 			return stream.Send(&agent.ResultResponse{File: data})
 		},
-		func(res interface{}) []byte {
+		func(res any) []byte {
 			return res.(*agent.ResultResponse).File
 		},
 	)
@@ -365,7 +365,7 @@ func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.Age
 		func(data []byte) error {
 			return stream.Send(&agent.AttestationResponse{File: data})
 		},
-		func(res interface{}) []byte {
+		func(res any) []byte {
 			return res.(*agent.AttestationResponse).File
 		},
 	)

agent/auth/auth.go

@@ -41,9 +41,9 @@ type Authenticator interface {
 }
 
 type service struct {
-	resultConsumers   []interface{}
-	datasetProviders  []interface{}
-	algorithmProvider interface{}
+	resultConsumers   []any
+	datasetProviders  []any
+	algorithmProvider any
 }
 
 func New(manifest agent.Computation) (Authenticator, error) {

agent/auth/mocks/authenticator.go

@@ -1,18 +1,33 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	context "context"
-
-	auth "github.com/ultravioletrs/cocos/agent/auth"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent/auth"
 )
 
+// NewAuthenticator creates a new instance of Authenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAuthenticator(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Authenticator {
+	mock := &Authenticator{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Authenticator is an autogenerated mock type for the Authenticator type
 type Authenticator struct {
 	mock.Mock
@@ -26,9 +41,9 @@ func (_m *Authenticator) EXPECT() *Authenticator_Expecter {
 	return &Authenticator_Expecter{mock: &_m.Mock}
 }
 
-// AuthenticateUser provides a mock function with given fields: ctx, role
-func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
-	ret := _m.Called(ctx, role)
+// AuthenticateUser provides a mock function for the type Authenticator
+func (_mock *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
+	ret := _mock.Called(ctx, role)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AuthenticateUser")
@@ -36,23 +51,21 @@ func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRol
 
 	var r0 context.Context
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
-		return rf(ctx, role)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
+		return returnFunc(ctx, role)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
-		r0 = rf(ctx, role)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
+		r0 = returnFunc(ctx, role)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(context.Context)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
-		r1 = rf(ctx, role)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
+		r1 = returnFunc(ctx, role)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -70,31 +83,28 @@ func (_e *Authenticator_Expecter) AuthenticateUser(ctx interface{}, role interfa
 
 func (_c *Authenticator_AuthenticateUser_Call) Run(run func(ctx context.Context, role auth.UserRole)) *Authenticator_AuthenticateUser_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(auth.UserRole))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 auth.UserRole
+		if args[1] != nil {
+			arg1 = args[1].(auth.UserRole)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Authenticator_AuthenticateUser_Call) Return(_a0 context.Context, _a1 error) *Authenticator_AuthenticateUser_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Authenticator_AuthenticateUser_Call) Return(context1 context.Context, err error) *Authenticator_AuthenticateUser_Call {
+	_c.Call.Return(context1, err)
 	return _c
 }
 
-func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(context.Context, auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
+func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(ctx context.Context, role auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewAuthenticator creates a new instance of Authenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewAuthenticator(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Authenticator {
-	mock := &Authenticator{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/cvms/api/grpc/storage/mocks/storage.go

@@ -1,17 +1,32 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
 	mock "github.com/stretchr/testify/mock"
-	cvms "github.com/ultravioletrs/cocos/agent/cvms"
-
-	storage "github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
 )
 
+// NewStorage creates a new instance of Storage. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewStorage(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Storage {
+	mock := &Storage{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Storage is an autogenerated mock type for the Storage type
 type Storage struct {
 	mock.Mock
@@ -25,21 +40,20 @@ func (_m *Storage) EXPECT() *Storage_Expecter {
 	return &Storage_Expecter{mock: &_m.Mock}
 }
 
-// Add provides a mock function with given fields: msg
-func (_m *Storage) Add(msg *cvms.ClientStreamMessage) error {
-	ret := _m.Called(msg)
+// Add provides a mock function for the type Storage
+func (_mock *Storage) Add(msg *cvms.ClientStreamMessage) error {
+	ret := _mock.Called(msg)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Add")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
-		r0 = rf(msg)
+	if returnFunc, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
+		r0 = returnFunc(msg)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -56,36 +70,41 @@ func (_e *Storage_Expecter) Add(msg interface{}) *Storage_Add_Call {
 
 func (_c *Storage_Add_Call) Run(run func(msg *cvms.ClientStreamMessage)) *Storage_Add_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(*cvms.ClientStreamMessage))
+		var arg0 *cvms.ClientStreamMessage
+		if args[0] != nil {
+			arg0 = args[0].(*cvms.ClientStreamMessage)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Storage_Add_Call) Return(_a0 error) *Storage_Add_Call {
-	_c.Call.Return(_a0)
+func (_c *Storage_Add_Call) Return(err error) *Storage_Add_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Storage_Add_Call) RunAndReturn(run func(*cvms.ClientStreamMessage) error) *Storage_Add_Call {
+func (_c *Storage_Add_Call) RunAndReturn(run func(msg *cvms.ClientStreamMessage) error) *Storage_Add_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Clear provides a mock function with no fields
-func (_m *Storage) Clear() error {
-	ret := _m.Called()
+// Clear provides a mock function for the type Storage
+func (_mock *Storage) Clear() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Clear")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -106,8 +125,8 @@ func (_c *Storage_Clear_Call) Run(run func()) *Storage_Clear_Call {
 	return _c
 }
 
-func (_c *Storage_Clear_Call) Return(_a0 error) *Storage_Clear_Call {
-	_c.Call.Return(_a0)
+func (_c *Storage_Clear_Call) Return(err error) *Storage_Clear_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -116,9 +135,9 @@ func (_c *Storage_Clear_Call) RunAndReturn(run func() error) *Storage_Clear_Call
 	return _c
 }
 
-// Load provides a mock function with no fields
-func (_m *Storage) Load() ([]storage.Message, error) {
-	ret := _m.Called()
+// Load provides a mock function for the type Storage
+func (_mock *Storage) Load() ([]storage.Message, error) {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -126,23 +145,21 @@ func (_m *Storage) Load() ([]storage.Message, error) {
 
 	var r0 []storage.Message
 	var r1 error
-	if rf, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
-		return rf()
+	if returnFunc, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
+		return returnFunc()
 	}
-	if rf, ok := ret.Get(0).(func() []storage.Message); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() []storage.Message); ok {
+		r0 = returnFunc()
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]storage.Message)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -163,8 +180,8 @@ func (_c *Storage_Load_Call) Run(run func()) *Storage_Load_Call {
 	return _c
 }
 
-func (_c *Storage_Load_Call) Return(_a0 []storage.Message, _a1 error) *Storage_Load_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Storage_Load_Call) Return(messages []storage.Message, err error) *Storage_Load_Call {
+	_c.Call.Return(messages, err)
 	return _c
 }
 
@@ -173,21 +190,20 @@ func (_c *Storage_Load_Call) RunAndReturn(run func() ([]storage.Message, error))
 	return _c
 }
 
-// Save provides a mock function with given fields: messages
-func (_m *Storage) Save(messages []storage.Message) error {
-	ret := _m.Called(messages)
+// Save provides a mock function for the type Storage
+func (_mock *Storage) Save(messages []storage.Message) error {
+	ret := _mock.Called(messages)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Save")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]storage.Message) error); ok {
-		r0 = rf(messages)
+	if returnFunc, ok := ret.Get(0).(func([]storage.Message) error); ok {
+		r0 = returnFunc(messages)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -204,31 +220,23 @@ func (_e *Storage_Expecter) Save(messages interface{}) *Storage_Save_Call {
 
 func (_c *Storage_Save_Call) Run(run func(messages []storage.Message)) *Storage_Save_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]storage.Message))
+		var arg0 []storage.Message
+		if args[0] != nil {
+			arg0 = args[0].([]storage.Message)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Storage_Save_Call) Return(_a0 error) *Storage_Save_Call {
-	_c.Call.Return(_a0)
+func (_c *Storage_Save_Call) Return(err error) *Storage_Save_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Storage_Save_Call) RunAndReturn(run func([]storage.Message) error) *Storage_Save_Call {
+func (_c *Storage_Save_Call) RunAndReturn(run func(messages []storage.Message) error) *Storage_Save_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewStorage creates a new instance of Storage. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewStorage(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Storage {
-	mock := &Storage{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/cvms/cvms.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: agent/cvms/cvms.proto
 
@@ -1263,177 +1263,94 @@ func (x *AzureAttestationToken) GetCertSerialNumber() string {
 
 var File_agent_cvms_cvms_proto protoreflect.FileDescriptor
 
-var file_agent_cvms_cvms_proto_rawDesc = string([]byte{
-	0x0a, 0x15, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x63, 0x76, 0x6d, 0x73, 0x2f, 0x63, 0x76, 0x6d,
-	0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x04, 0x63, 0x76, 0x6d, 0x73, 0x1a, 0x1f, 0x67,
-	0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74,
-	0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x1f,
-	0x0a, 0x0d, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x12,
-	0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22,
-	0x35, 0x0a, 0x0d, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73,
-	0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64,
-	0x12, 0x14, 0x0a, 0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x22, 0x38, 0x0a, 0x0f, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f,
-	0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6f, 0x6d,
-	0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64,
-	0x22, 0x5a, 0x0a, 0x17, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74,
-	0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x25, 0x0a, 0x0e, 0x63,
-	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
-	0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x02, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x4a, 0x0a, 0x0b,
-	0x52, 0x75, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x25, 0x0a, 0x0e, 0x63,
-	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
-	0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0xde, 0x01, 0x0a, 0x0a, 0x41, 0x67, 0x65,
-	0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x76, 0x65, 0x6e, 0x74,
-	0x5f, 0x74, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x65, 0x76, 0x65,
-	0x6e, 0x74, 0x54, 0x79, 0x70, 0x65, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74,
-	0x61, 0x6d, 0x70, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67,
-	0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65,
-	0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70,
-	0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f,
-	0x69, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74,
-	0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69,
-	0x6c, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c,
-	0x73, 0x12, 0x1e, 0x0a, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x18,
-	0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f,
-	0x72, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x9b, 0x01, 0x0a, 0x08, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67,
-	0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65,
-	0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f,
-	0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74,
-	0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c,
-	0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x38, 0x0a,
-	0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b,
-	0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62,
-	0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69,
-	0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x22, 0xe8, 0x03, 0x0a, 0x13, 0x43, 0x6c, 0x69, 0x65,
-	0x6e, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12,
-	0x2d, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c,
-	0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x12, 0x33,
-	0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x18, 0x02, 0x20,
-	0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74,
-	0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76,
-	0x65, 0x6e, 0x74, 0x12, 0x2c, 0x0a, 0x07, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x65, 0x73, 0x18, 0x03,
-	0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x52, 0x75, 0x6e, 0x52,
-	0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65,
-	0x73, 0x12, 0x4f, 0x0a, 0x12, 0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1d, 0x2e,
-	0x63, 0x76, 0x6d, 0x73, 0x2e, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x12,
-	0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52,
-	0x65, 0x73, 0x12, 0x3b, 0x0a, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65,
-	0x52, 0x65, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x13, 0x2e, 0x63, 0x76, 0x6d, 0x73,
-	0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x48, 0x00,
-	0x52, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x12,
-	0x51, 0x0a, 0x15, 0x76, 0x54, 0x50, 0x4d, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69,
-	0x6f, 0x6e, 0x52, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19,
-	0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f,
-	0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x15, 0x76, 0x54, 0x50,
-	0x4d, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x70, 0x6f,
-	0x72, 0x74, 0x12, 0x53, 0x0a, 0x15, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x18, 0x07, 0x20, 0x01, 0x28,
-	0x0b, 0x32, 0x1b, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41, 0x74,
-	0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x48, 0x00,
-	0x52, 0x15, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69,
-	0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61,
-	0x67, 0x65, 0x22, 0xca, 0x02, 0x0a, 0x13, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x53, 0x74, 0x72,
-	0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x38, 0x0a, 0x0c, 0x72, 0x75,
-	0x6e, 0x52, 0x65, 0x71, 0x43, 0x68, 0x75, 0x6e, 0x6b, 0x73, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b,
-	0x32, 0x12, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x43, 0x68,
-	0x75, 0x6e, 0x6b, 0x73, 0x48, 0x00, 0x52, 0x0c, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x43, 0x68,
-	0x75, 0x6e, 0x6b, 0x73, 0x12, 0x31, 0x0a, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x18, 0x02,
-	0x20, 0x01, 0x28, 0x0b, 0x32, 0x17, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x43, 0x6f, 0x6d, 0x70,
-	0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x48, 0x00, 0x52,
-	0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x12, 0x41, 0x0a, 0x0f, 0x73, 0x74, 0x6f, 0x70, 0x43,
-	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b,
-	0x32, 0x15, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70,
-	0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x48, 0x00, 0x52, 0x0f, 0x73, 0x74, 0x6f, 0x70, 0x43,
-	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x3b, 0x0a, 0x0d, 0x61, 0x67,
-	0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x18, 0x04, 0x20, 0x01, 0x28,
-	0x0b, 0x32, 0x13, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74,
-	0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x48, 0x00, 0x52, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x53,
-	0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x12, 0x3b, 0x0a, 0x0d, 0x64, 0x69, 0x73, 0x63, 0x6f,
-	0x6e, 0x6e, 0x65, 0x63, 0x74, 0x52, 0x65, 0x71, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x13,
-	0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x44, 0x69, 0x73, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74,
-	0x52, 0x65, 0x71, 0x48, 0x00, 0x52, 0x0d, 0x64, 0x69, 0x73, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63,
-	0x74, 0x52, 0x65, 0x71, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22,
-	0x1f, 0x0a, 0x0d, 0x44, 0x69, 0x73, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x52, 0x65, 0x71,
-	0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64,
-	0x22, 0x4b, 0x0a, 0x0c, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x43, 0x68, 0x75, 0x6e, 0x6b, 0x73,
-	0x12, 0x12, 0x0a, 0x04, 0x64, 0x61, 0x74, 0x61, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04,
-	0x64, 0x61, 0x74, 0x61, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09,
-	0x52, 0x02, 0x69, 0x64, 0x12, 0x17, 0x0a, 0x07, 0x69, 0x73, 0x5f, 0x6c, 0x61, 0x73, 0x74, 0x18,
-	0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x06, 0x69, 0x73, 0x4c, 0x61, 0x73, 0x74, 0x22, 0xaa, 0x02,
-	0x0a, 0x11, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x75, 0x6e,
-	0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x20, 0x0a, 0x0b, 0x64, 0x65, 0x73, 0x63, 0x72,
-	0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0b, 0x64, 0x65,
-	0x73, 0x63, 0x72, 0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x29, 0x0a, 0x08, 0x64, 0x61, 0x74,
-	0x61, 0x73, 0x65, 0x74, 0x73, 0x18, 0x04, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x0d, 0x2e, 0x63, 0x76,
-	0x6d, 0x73, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x52, 0x08, 0x64, 0x61, 0x74, 0x61,
-	0x73, 0x65, 0x74, 0x73, 0x12, 0x2d, 0x0a, 0x09, 0x61, 0x6c, 0x67, 0x6f, 0x72, 0x69, 0x74, 0x68,
-	0x6d, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0f, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41,
-	0x6c, 0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x52, 0x09, 0x61, 0x6c, 0x67, 0x6f, 0x72, 0x69,
-	0x74, 0x68, 0x6d, 0x12, 0x3f, 0x0a, 0x10, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x5f, 0x63, 0x6f,
-	0x6e, 0x73, 0x75, 0x6d, 0x65, 0x72, 0x73, 0x18, 0x06, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x14, 0x2e,
-	0x63, 0x76, 0x6d, 0x73, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x43, 0x6f, 0x6e, 0x73, 0x75,
-	0x6d, 0x65, 0x72, 0x52, 0x0f, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x43, 0x6f, 0x6e, 0x73, 0x75,
-	0x6d, 0x65, 0x72, 0x73, 0x12, 0x34, 0x0a, 0x0c, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x6f,
-	0x6e, 0x66, 0x69, 0x67, 0x18, 0x07, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x63, 0x76, 0x6d,
-	0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x52, 0x0b, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x22, 0x2a, 0x0a, 0x0e, 0x52, 0x65,
-	0x73, 0x75, 0x6c, 0x74, 0x43, 0x6f, 0x6e, 0x73, 0x75, 0x6d, 0x65, 0x72, 0x12, 0x18, 0x0a, 0x07,
-	0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x75,
-	0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x22, 0x53, 0x0a, 0x07, 0x44, 0x61, 0x74, 0x61, 0x73, 0x65,
-	0x74, 0x12, 0x12, 0x0a, 0x04, 0x68, 0x61, 0x73, 0x68, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52,
-	0x04, 0x68, 0x61, 0x73, 0x68, 0x12, 0x18, 0x0a, 0x07, 0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79,
-	0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x12,
-	0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x39, 0x0a, 0x09, 0x41,
-	0x6c, 0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x12, 0x0a, 0x04, 0x68, 0x61, 0x73, 0x68,
-	0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x68, 0x61, 0x73, 0x68, 0x12, 0x18, 0x0a, 0x07,
-	0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x75,
-	0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x22, 0xe5, 0x01, 0x0a, 0x0b, 0x41, 0x67, 0x65, 0x6e, 0x74,
-	0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x12, 0x1b, 0x0a, 0x09, 0x63, 0x65,
-	0x72, 0x74, 0x5f, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x63,
-	0x65, 0x72, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x19, 0x0a, 0x08, 0x6b, 0x65, 0x79, 0x5f, 0x66,
-	0x69, 0x6c, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x6b, 0x65, 0x79, 0x46, 0x69,
-	0x6c, 0x65, 0x12, 0x24, 0x0a, 0x0e, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x61, 0x5f,
-	0x66, 0x69, 0x6c, 0x65, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x63, 0x6c, 0x69, 0x65,
-	0x6e, 0x74, 0x43, 0x61, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x24, 0x0a, 0x0e, 0x73, 0x65, 0x72, 0x76,
-	0x65, 0x72, 0x5f, 0x63, 0x61, 0x5f, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09,
-	0x52, 0x0c, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x61, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x1b,
-	0x0a, 0x09, 0x6c, 0x6f, 0x67, 0x5f, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x06, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x08, 0x6c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x21, 0x0a, 0x0c, 0x61,
-	0x74, 0x74, 0x65, 0x73, 0x74, 0x65, 0x64, 0x5f, 0x74, 0x6c, 0x73, 0x18, 0x07, 0x20, 0x01, 0x28,
-	0x08, 0x52, 0x0b, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x65, 0x64, 0x54, 0x6c, 0x73, 0x22, 0x55,
-	0x0a, 0x13, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x12, 0x2a, 0x0a, 0x10, 0x63, 0x65, 0x72,
-	0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x18, 0x02, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x10, 0x63, 0x65, 0x72, 0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x4e,
-	0x75, 0x6d, 0x62, 0x65, 0x72, 0x22, 0x57, 0x0a, 0x15, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41, 0x74,
-	0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x12, 0x12,
-	0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69,
-	0x6c, 0x65, 0x12, 0x2a, 0x0a, 0x10, 0x63, 0x65, 0x72, 0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c,
-	0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x63, 0x65,
-	0x72, 0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x32, 0x50,
-	0x0a, 0x07, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x45, 0x0a, 0x07, 0x50, 0x72, 0x6f,
-	0x63, 0x65, 0x73, 0x73, 0x12, 0x19, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x43, 0x6c, 0x69, 0x65,
-	0x6e, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a,
-	0x19, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x53, 0x74, 0x72,
-	0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x00, 0x28, 0x01, 0x30, 0x01,
-	0x42, 0x08, 0x5a, 0x06, 0x2e, 0x2f, 0x63, 0x76, 0x6d, 0x73, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74,
-	0x6f, 0x33,
-})
+const file_agent_cvms_cvms_proto_rawDesc = "" +
+	"\n" +
+	"\x15agent/cvms/cvms.proto\x12\x04cvms\x1a\x1fgoogle/protobuf/timestamp.proto\"\x1f\n" +
+	"\rAgentStateReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\"5\n" +
+	"\rAgentStateRes\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\x12\x14\n" +
+	"\x05state\x18\x02 \x01(\tR\x05state\"8\n" +
+	"\x0fStopComputation\x12%\n" +
+	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\"Z\n" +
+	"\x17StopComputationResponse\x12%\n" +
+	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x18\n" +
+	"\amessage\x18\x02 \x01(\tR\amessage\"J\n" +
+	"\vRunResponse\x12%\n" +
+	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x14\n" +
+	"\x05error\x18\x02 \x01(\tR\x05error\"\xde\x01\n" +
+	"\n" +
+	"AgentEvent\x12\x1d\n" +
+	"\n" +
+	"event_type\x18\x01 \x01(\tR\teventType\x128\n" +
+	"\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12%\n" +
+	"\x0ecomputation_id\x18\x03 \x01(\tR\rcomputationId\x12\x18\n" +
+	"\adetails\x18\x04 \x01(\fR\adetails\x12\x1e\n" +
+	"\n" +
+	"originator\x18\x05 \x01(\tR\n" +
+	"originator\x12\x16\n" +
+	"\x06status\x18\x06 \x01(\tR\x06status\"\x9b\x01\n" +
+	"\bAgentLog\x12\x18\n" +
+	"\amessage\x18\x01 \x01(\tR\amessage\x12%\n" +
+	"\x0ecomputation_id\x18\x02 \x01(\tR\rcomputationId\x12\x14\n" +
+	"\x05level\x18\x03 \x01(\tR\x05level\x128\n" +
+	"\ttimestamp\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\xe8\x03\n" +
+	"\x13ClientStreamMessage\x12-\n" +
+	"\tagent_log\x18\x01 \x01(\v2\x0e.cvms.AgentLogH\x00R\bagentLog\x123\n" +
+	"\vagent_event\x18\x02 \x01(\v2\x10.cvms.AgentEventH\x00R\n" +
+	"agentEvent\x12,\n" +
+	"\arun_res\x18\x03 \x01(\v2\x11.cvms.RunResponseH\x00R\x06runRes\x12O\n" +
+	"\x12stopComputationRes\x18\x04 \x01(\v2\x1d.cvms.StopComputationResponseH\x00R\x12stopComputationRes\x12;\n" +
+	"\ragentStateRes\x18\x05 \x01(\v2\x13.cvms.AgentStateResH\x00R\ragentStateRes\x12Q\n" +
+	"\x15vTPMattestationReport\x18\x06 \x01(\v2\x19.cvms.AttestationResponseH\x00R\x15vTPMattestationReport\x12S\n" +
+	"\x15azureAttestationToken\x18\a \x01(\v2\x1b.cvms.azureAttestationTokenH\x00R\x15azureAttestationTokenB\t\n" +
+	"\amessage\"\xca\x02\n" +
+	"\x13ServerStreamMessage\x128\n" +
+	"\frunReqChunks\x18\x01 \x01(\v2\x12.cvms.RunReqChunksH\x00R\frunReqChunks\x121\n" +
+	"\x06runReq\x18\x02 \x01(\v2\x17.cvms.ComputationRunReqH\x00R\x06runReq\x12A\n" +
+	"\x0fstopComputation\x18\x03 \x01(\v2\x15.cvms.StopComputationH\x00R\x0fstopComputation\x12;\n" +
+	"\ragentStateReq\x18\x04 \x01(\v2\x13.cvms.AgentStateReqH\x00R\ragentStateReq\x12;\n" +
+	"\rdisconnectReq\x18\x05 \x01(\v2\x13.cvms.DisconnectReqH\x00R\rdisconnectReqB\t\n" +
+	"\amessage\"\x1f\n" +
+	"\rDisconnectReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\"K\n" +
+	"\fRunReqChunks\x12\x12\n" +
+	"\x04data\x18\x01 \x01(\fR\x04data\x12\x0e\n" +
+	"\x02id\x18\x02 \x01(\tR\x02id\x12\x17\n" +
+	"\ais_last\x18\x03 \x01(\bR\x06isLast\"\xaa\x02\n" +
+	"\x11ComputationRunReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\x12\x12\n" +
+	"\x04name\x18\x02 \x01(\tR\x04name\x12 \n" +
+	"\vdescription\x18\x03 \x01(\tR\vdescription\x12)\n" +
+	"\bdatasets\x18\x04 \x03(\v2\r.cvms.DatasetR\bdatasets\x12-\n" +
+	"\talgorithm\x18\x05 \x01(\v2\x0f.cvms.AlgorithmR\talgorithm\x12?\n" +
+	"\x10result_consumers\x18\x06 \x03(\v2\x14.cvms.ResultConsumerR\x0fresultConsumers\x124\n" +
+	"\fagent_config\x18\a \x01(\v2\x11.cvms.AgentConfigR\vagentConfig\"*\n" +
+	"\x0eResultConsumer\x12\x18\n" +
+	"\auserKey\x18\x01 \x01(\fR\auserKey\"S\n" +
+	"\aDataset\x12\x12\n" +
+	"\x04hash\x18\x01 \x01(\fR\x04hash\x12\x18\n" +
+	"\auserKey\x18\x02 \x01(\fR\auserKey\x12\x1a\n" +
+	"\bfilename\x18\x03 \x01(\tR\bfilename\"9\n" +
+	"\tAlgorithm\x12\x12\n" +
+	"\x04hash\x18\x01 \x01(\fR\x04hash\x12\x18\n" +
+	"\auserKey\x18\x02 \x01(\fR\auserKey\"\xe5\x01\n" +
+	"\vAgentConfig\x12\x12\n" +
+	"\x04port\x18\x01 \x01(\tR\x04port\x12\x1b\n" +
+	"\tcert_file\x18\x02 \x01(\tR\bcertFile\x12\x19\n" +
+	"\bkey_file\x18\x03 \x01(\tR\akeyFile\x12$\n" +
+	"\x0eclient_ca_file\x18\x04 \x01(\tR\fclientCaFile\x12$\n" +
+	"\x0eserver_ca_file\x18\x05 \x01(\tR\fserverCaFile\x12\x1b\n" +
+	"\tlog_level\x18\x06 \x01(\tR\blogLevel\x12!\n" +
+	"\fattested_tls\x18\a \x01(\bR\vattestedTls\"U\n" +
+	"\x13AttestationResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\x12*\n" +
+	"\x10certSerialNumber\x18\x02 \x01(\tR\x10certSerialNumber\"W\n" +
+	"\x15azureAttestationToken\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\x12*\n" +
+	"\x10certSerialNumber\x18\x02 \x01(\tR\x10certSerialNumber2P\n" +
+	"\aService\x12E\n" +
+	"\aProcess\x12\x19.cvms.ClientStreamMessage\x1a\x19.cvms.ServerStreamMessage\"\x00(\x010\x01B\bZ\x06./cvmsb\x06proto3"
 
 var (
 	file_agent_cvms_cvms_proto_rawDescOnce sync.Once

agent/cvms/server/mocks/agentserver.go

@@ -1,15 +1,31 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
 	mock "github.com/stretchr/testify/mock"
-	agent "github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent"
 )
 
+// NewAgentServer creates a new instance of AgentServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAgentServer(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentServer {
+	mock := &AgentServer{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // AgentServer is an autogenerated mock type for the AgentServer type
 type AgentServer struct {
 	mock.Mock
@@ -23,21 +39,20 @@ func (_m *AgentServer) EXPECT() *AgentServer_Expecter {
 	return &AgentServer_Expecter{mock: &_m.Mock}
 }
 
-// Start provides a mock function with given fields: cfg, cmp
-func (_m *AgentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
-	ret := _m.Called(cfg, cmp)
+// Start provides a mock function for the type AgentServer
+func (_mock *AgentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
+	ret := _mock.Called(cfg, cmp)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Start")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
-		r0 = rf(cfg, cmp)
+	if returnFunc, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
+		r0 = returnFunc(cfg, cmp)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -55,36 +70,46 @@ func (_e *AgentServer_Expecter) Start(cfg interface{}, cmp interface{}) *AgentSe
 
 func (_c *AgentServer_Start_Call) Run(run func(cfg agent.AgentConfig, cmp agent.Computation)) *AgentServer_Start_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(agent.AgentConfig), args[1].(agent.Computation))
+		var arg0 agent.AgentConfig
+		if args[0] != nil {
+			arg0 = args[0].(agent.AgentConfig)
+		}
+		var arg1 agent.Computation
+		if args[1] != nil {
+			arg1 = args[1].(agent.Computation)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *AgentServer_Start_Call) Return(_a0 error) *AgentServer_Start_Call {
-	_c.Call.Return(_a0)
+func (_c *AgentServer_Start_Call) Return(err error) *AgentServer_Start_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *AgentServer_Start_Call) RunAndReturn(run func(agent.AgentConfig, agent.Computation) error) *AgentServer_Start_Call {
+func (_c *AgentServer_Start_Call) RunAndReturn(run func(cfg agent.AgentConfig, cmp agent.Computation) error) *AgentServer_Start_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Stop provides a mock function with no fields
-func (_m *AgentServer) Stop() error {
-	ret := _m.Called()
+// Stop provides a mock function for the type AgentServer
+func (_mock *AgentServer) Stop() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Stop")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -105,8 +130,8 @@ func (_c *AgentServer_Stop_Call) Run(run func()) *AgentServer_Stop_Call {
 	return _c
 }
 
-func (_c *AgentServer_Stop_Call) Return(_a0 error) *AgentServer_Stop_Call {
-	_c.Call.Return(_a0)
+func (_c *AgentServer_Stop_Call) Return(err error) *AgentServer_Stop_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -114,17 +139,3 @@ func (_c *AgentServer_Stop_Call) RunAndReturn(run func() error) *AgentServer_Sto
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewAgentServer creates a new instance of AgentServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewAgentServer(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentServer {
-	mock := &AgentServer{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/events/events.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: agent/events/events.proto
 
@@ -261,46 +261,32 @@ func (*EventsLogs_AgentEvent) isEventsLogs_Message() {}
 
 var File_agent_events_events_proto protoreflect.FileDescriptor
 
-var file_agent_events_events_proto_rawDesc = string([]byte{
-	0x0a, 0x19, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2f, 0x65,
-	0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x06, 0x65, 0x76, 0x65,
-	0x6e, 0x74, 0x73, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74,
-	0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70,
-	0x72, 0x6f, 0x74, 0x6f, 0x22, 0xde, 0x01, 0x0a, 0x0a, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76,
-	0x65, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x5f, 0x74, 0x79, 0x70,
-	0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x54, 0x79,
-	0x70, 0x65, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18,
-	0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70,
-	0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d,
-	0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x12, 0x25, 0x0a, 0x0e,
-	0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x03,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
-	0x6e, 0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x18, 0x04,
-	0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x12, 0x1e, 0x0a,
-	0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x16, 0x0a,
-	0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73,
-	0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x9b, 0x01, 0x0a, 0x08, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c,
-	0x6f, 0x67, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x25, 0x0a, 0x0e,
-	0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x02,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
-	0x6e, 0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x03, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d,
-	0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67,
-	0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54,
-	0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74,
-	0x61, 0x6d, 0x70, 0x22, 0x7f, 0x0a, 0x0a, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x4c, 0x6f, 0x67,
-	0x73, 0x12, 0x2f, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c,
-	0x6f, 0x67, 0x12, 0x35, 0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e,
-	0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
-	0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73,
-	0x73, 0x61, 0x67, 0x65, 0x42, 0x0a, 0x5a, 0x08, 0x2e, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
-	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-})
+const file_agent_events_events_proto_rawDesc = "" +
+	"\n" +
+	"\x19agent/events/events.proto\x12\x06events\x1a\x1fgoogle/protobuf/timestamp.proto\"\xde\x01\n" +
+	"\n" +
+	"AgentEvent\x12\x1d\n" +
+	"\n" +
+	"event_type\x18\x01 \x01(\tR\teventType\x128\n" +
+	"\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12%\n" +
+	"\x0ecomputation_id\x18\x03 \x01(\tR\rcomputationId\x12\x18\n" +
+	"\adetails\x18\x04 \x01(\fR\adetails\x12\x1e\n" +
+	"\n" +
+	"originator\x18\x05 \x01(\tR\n" +
+	"originator\x12\x16\n" +
+	"\x06status\x18\x06 \x01(\tR\x06status\"\x9b\x01\n" +
+	"\bAgentLog\x12\x18\n" +
+	"\amessage\x18\x01 \x01(\tR\amessage\x12%\n" +
+	"\x0ecomputation_id\x18\x02 \x01(\tR\rcomputationId\x12\x14\n" +
+	"\x05level\x18\x03 \x01(\tR\x05level\x128\n" +
+	"\ttimestamp\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\x7f\n" +
+	"\n" +
+	"EventsLogs\x12/\n" +
+	"\tagent_log\x18\x01 \x01(\v2\x10.events.AgentLogH\x00R\bagentLog\x125\n" +
+	"\vagent_event\x18\x02 \x01(\v2\x12.events.AgentEventH\x00R\n" +
+	"agentEventB\t\n" +
+	"\amessageB\n" +
+	"Z\b./eventsb\x06proto3"
 
 var (
 	file_agent_events_events_proto_rawDescOnce sync.Once

agent/events/mocks/service.go

@@ -1,16 +1,32 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	json "encoding/json"
+	"encoding/json"
 
 	mock "github.com/stretchr/testify/mock"
 )
 
+// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewService(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Service is an autogenerated mock type for the Service type
 type Service struct {
 	mock.Mock
@@ -24,9 +40,10 @@ func (_m *Service) EXPECT() *Service_Expecter {
 	return &Service_Expecter{mock: &_m.Mock}
 }
 
-// SendEvent provides a mock function with given fields: cmpID, event, status, details
-func (_m *Service) SendEvent(cmpID string, event string, status string, details json.RawMessage) {
-	_m.Called(cmpID, event, status, details)
+// SendEvent provides a mock function for the type Service
+func (_mock *Service) SendEvent(cmpID string, event string, status string, details json.RawMessage) {
+	_mock.Called(cmpID, event, status, details)
+	return
 }
 
 // Service_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
@@ -45,7 +62,28 @@ func (_e *Service_Expecter) SendEvent(cmpID interface{}, event interface{}, stat
 
 func (_c *Service_SendEvent_Call) Run(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string), args[2].(string), args[3].(json.RawMessage))
+		var arg0 string
+		if args[0] != nil {
+			arg0 = args[0].(string)
+		}
+		var arg1 string
+		if args[1] != nil {
+			arg1 = args[1].(string)
+		}
+		var arg2 string
+		if args[2] != nil {
+			arg2 = args[2].(string)
+		}
+		var arg3 json.RawMessage
+		if args[3] != nil {
+			arg3 = args[3].(json.RawMessage)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+			arg3,
+		)
 	})
 	return _c
 }
@@ -55,21 +93,7 @@ func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
 	return _c
 }
 
-func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, string, json.RawMessage)) *Service_SendEvent_Call {
+func (_c *Service_SendEvent_Call) RunAndReturn(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
 	_c.Run(run)
 	return _c
 }
-
-// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewService(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Service {
-	mock := &Service{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agent_grpc_algo.go

@@ -1,434 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	metadata "google.golang.org/grpc/metadata"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
-type AgentService_AlgoClient[Req interface{}, Res interface{}] struct {
-	mock.Mock
-}
-
-type AgentService_AlgoClient_Expecter[Req interface{}, Res interface{}] struct {
-	mock *mock.Mock
-}
-
-func (_m *AgentService_AlgoClient[Req, Res]) EXPECT() *AgentService_AlgoClient_Expecter[Req, Res] {
-	return &AgentService_AlgoClient_Expecter[Req, Res]{mock: &_m.Mock}
-}
-
-// CloseAndRecv provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) CloseAndRecv() (*agent.AlgoResponse, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseAndRecv")
-	}
-
-	var r0 *agent.AlgoResponse
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*agent.AlgoResponse)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_AlgoClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
-type AgentService_AlgoClient_CloseAndRecv_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseAndRecv is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) CloseAndRecv() *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	return &AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]{Call: _e.mock.On("CloseAndRecv")}
-}
-
-func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) Return(_a0 *agent.AlgoResponse, _a1 error) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) RunAndReturn(run func() (*agent.AlgoResponse, error)) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// CloseSend provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) CloseSend() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseSend")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
-type AgentService_AlgoClient_CloseSend_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseSend is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) CloseSend() *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	return &AgentService_AlgoClient_CloseSend_Call[Req, Res]{Call: _e.mock.On("CloseSend")}
-}
-
-func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) RunAndReturn(run func() error) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Context provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) Context() context.Context {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Context")
-	}
-
-	var r0 context.Context
-	if rf, ok := ret.Get(0).(func() context.Context); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(context.Context)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
-type AgentService_AlgoClient_Context_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Context is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Context() *AgentService_AlgoClient_Context_Call[Req, Res] {
-	return &AgentService_AlgoClient_Context_Call[Req, Res]{Call: _e.mock.On("Context")}
-}
-
-func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Context_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) Return(_a0 context.Context) *AgentService_AlgoClient_Context_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) RunAndReturn(run func() context.Context) *AgentService_AlgoClient_Context_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Header provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) Header() (metadata.MD, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Header")
-	}
-
-	var r0 metadata.MD
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_AlgoClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
-type AgentService_AlgoClient_Header_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Header is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Header() *AgentService_AlgoClient_Header_Call[Req, Res] {
-	return &AgentService_AlgoClient_Header_Call[Req, Res]{Call: _e.mock.On("Header")}
-}
-
-func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Header_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_AlgoClient_Header_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_AlgoClient_Header_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// RecvMsg provides a mock function with given fields: m
-func (_m *AgentService_AlgoClient[Req, Res]) RecvMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for RecvMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
-type AgentService_AlgoClient_RecvMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// RecvMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) RecvMsg(m interface{}) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	return &AgentService_AlgoClient_RecvMsg_Call[Req, Res]{Call: _e.mock.On("RecvMsg", m)}
-}
-
-func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Send provides a mock function with given fields: _a0
-func (_m *AgentService_AlgoClient[Req, Res]) Send(_a0 *agent.AlgoRequest) error {
-	ret := _m.Called(_a0)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Send")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
-		r0 = rf(_a0)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
-type AgentService_AlgoClient_Send_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Send is a helper method to define mock.On call
-//   - _a0 *agent.AlgoRequest
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Send(_a0 interface{}) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	return &AgentService_AlgoClient_Send_Call[Req, Res]{Call: _e.mock.On("Send", _a0)}
-}
-
-func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) Run(run func(_a0 *agent.AlgoRequest)) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(*agent.AlgoRequest))
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) RunAndReturn(run func(*agent.AlgoRequest) error) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// SendMsg provides a mock function with given fields: m
-func (_m *AgentService_AlgoClient[Req, Res]) SendMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
-type AgentService_AlgoClient_SendMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// SendMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) SendMsg(m interface{}) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	return &AgentService_AlgoClient_SendMsg_Call[Req, Res]{Call: _e.mock.On("SendMsg", m)}
-}
-
-func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Trailer provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) Trailer() metadata.MD {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Trailer")
-	}
-
-	var r0 metadata.MD
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
-type AgentService_AlgoClient_Trailer_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Trailer is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Trailer() *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	return &AgentService_AlgoClient_Trailer_Call[Req, Res]{Call: _e.mock.On("Trailer")}
-}
-
-func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) Return(_a0 metadata.MD) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) RunAndReturn(run func() metadata.MD) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewAgentService_AlgoClient creates a new instance of AgentService_AlgoClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewAgentService_AlgoClient[Req interface{}, Res interface{}](t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentService_AlgoClient[Req, Res] {
-	mock := &AgentService_AlgoClient[Req, Res]{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agent_grpc_data.go

@@ -1,434 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	metadata "google.golang.org/grpc/metadata"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
-type AgentService_DataClient[Req interface{}, Res interface{}] struct {
-	mock.Mock
-}
-
-type AgentService_DataClient_Expecter[Req interface{}, Res interface{}] struct {
-	mock *mock.Mock
-}
-
-func (_m *AgentService_DataClient[Req, Res]) EXPECT() *AgentService_DataClient_Expecter[Req, Res] {
-	return &AgentService_DataClient_Expecter[Req, Res]{mock: &_m.Mock}
-}
-
-// CloseAndRecv provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) CloseAndRecv() (*agent.DataResponse, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseAndRecv")
-	}
-
-	var r0 *agent.DataResponse
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *agent.DataResponse); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*agent.DataResponse)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_DataClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
-type AgentService_DataClient_CloseAndRecv_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseAndRecv is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) CloseAndRecv() *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	return &AgentService_DataClient_CloseAndRecv_Call[Req, Res]{Call: _e.mock.On("CloseAndRecv")}
-}
-
-func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) Run(run func()) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) Return(_a0 *agent.DataResponse, _a1 error) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) RunAndReturn(run func() (*agent.DataResponse, error)) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// CloseSend provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) CloseSend() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseSend")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
-type AgentService_DataClient_CloseSend_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseSend is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) CloseSend() *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	return &AgentService_DataClient_CloseSend_Call[Req, Res]{Call: _e.mock.On("CloseSend")}
-}
-
-func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) Run(run func()) *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) RunAndReturn(run func() error) *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Context provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) Context() context.Context {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Context")
-	}
-
-	var r0 context.Context
-	if rf, ok := ret.Get(0).(func() context.Context); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(context.Context)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
-type AgentService_DataClient_Context_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Context is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Context() *AgentService_DataClient_Context_Call[Req, Res] {
-	return &AgentService_DataClient_Context_Call[Req, Res]{Call: _e.mock.On("Context")}
-}
-
-func (_c *AgentService_DataClient_Context_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Context_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Context_Call[Req, Res]) Return(_a0 context.Context) *AgentService_DataClient_Context_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Context_Call[Req, Res]) RunAndReturn(run func() context.Context) *AgentService_DataClient_Context_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Header provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) Header() (metadata.MD, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Header")
-	}
-
-	var r0 metadata.MD
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_DataClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
-type AgentService_DataClient_Header_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Header is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Header() *AgentService_DataClient_Header_Call[Req, Res] {
-	return &AgentService_DataClient_Header_Call[Req, Res]{Call: _e.mock.On("Header")}
-}
-
-func (_c *AgentService_DataClient_Header_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Header_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Header_Call[Req, Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_DataClient_Header_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Header_Call[Req, Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_DataClient_Header_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// RecvMsg provides a mock function with given fields: m
-func (_m *AgentService_DataClient[Req, Res]) RecvMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for RecvMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
-type AgentService_DataClient_RecvMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// RecvMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_DataClient_Expecter[Req, Res]) RecvMsg(m interface{}) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	return &AgentService_DataClient_RecvMsg_Call[Req, Res]{Call: _e.mock.On("RecvMsg", m)}
-}
-
-func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Send provides a mock function with given fields: _a0
-func (_m *AgentService_DataClient[Req, Res]) Send(_a0 *agent.DataRequest) error {
-	ret := _m.Called(_a0)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Send")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
-		r0 = rf(_a0)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
-type AgentService_DataClient_Send_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Send is a helper method to define mock.On call
-//   - _a0 *agent.DataRequest
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Send(_a0 interface{}) *AgentService_DataClient_Send_Call[Req, Res] {
-	return &AgentService_DataClient_Send_Call[Req, Res]{Call: _e.mock.On("Send", _a0)}
-}
-
-func (_c *AgentService_DataClient_Send_Call[Req, Res]) Run(run func(_a0 *agent.DataRequest)) *AgentService_DataClient_Send_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(*agent.DataRequest))
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Send_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_Send_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Send_Call[Req, Res]) RunAndReturn(run func(*agent.DataRequest) error) *AgentService_DataClient_Send_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// SendMsg provides a mock function with given fields: m
-func (_m *AgentService_DataClient[Req, Res]) SendMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
-type AgentService_DataClient_SendMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// SendMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_DataClient_Expecter[Req, Res]) SendMsg(m interface{}) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	return &AgentService_DataClient_SendMsg_Call[Req, Res]{Call: _e.mock.On("SendMsg", m)}
-}
-
-func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Trailer provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) Trailer() metadata.MD {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Trailer")
-	}
-
-	var r0 metadata.MD
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
-type AgentService_DataClient_Trailer_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Trailer is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Trailer() *AgentService_DataClient_Trailer_Call[Req, Res] {
-	return &AgentService_DataClient_Trailer_Call[Req, Res]{Call: _e.mock.On("Trailer")}
-}
-
-func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Trailer_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) Return(_a0 metadata.MD) *AgentService_DataClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) RunAndReturn(run func() metadata.MD) *AgentService_DataClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewAgentService_DataClient creates a new instance of AgentService_DataClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewAgentService_DataClient[Req interface{}, Res interface{}](t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentService_DataClient[Req, Res] {
-	mock := &AgentService_DataClient[Req, Res]{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agent_grpc_ima.go

@@ -1,388 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	metadata "google.golang.org/grpc/metadata"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// AgentService_IMAMeasurementsClient is an autogenerated mock type for the AgentService_IMAMeasurementsClient type
-type AgentService_IMAMeasurementsClient[Res interface{}] struct {
-	mock.Mock
-}
-
-type AgentService_IMAMeasurementsClient_Expecter[Res interface{}] struct {
-	mock *mock.Mock
-}
-
-func (_m *AgentService_IMAMeasurementsClient[Res]) EXPECT() *AgentService_IMAMeasurementsClient_Expecter[Res] {
-	return &AgentService_IMAMeasurementsClient_Expecter[Res]{mock: &_m.Mock}
-}
-
-// CloseSend provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) CloseSend() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseSend")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
-type AgentService_IMAMeasurementsClient_CloseSend_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseSend is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) CloseSend() *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_CloseSend_Call[Res]{Call: _e.mock.On("CloseSend")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call[Res]) Return(_a0 error) *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call[Res]) RunAndReturn(run func() error) *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Context provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Context() context.Context {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Context")
-	}
-
-	var r0 context.Context
-	if rf, ok := ret.Get(0).(func() context.Context); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(context.Context)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
-type AgentService_IMAMeasurementsClient_Context_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Context is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Context() *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Context_Call[Res]{Call: _e.mock.On("Context")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Context_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Context_Call[Res]) Return(_a0 context.Context) *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Context_Call[Res]) RunAndReturn(run func() context.Context) *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Header provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Header() (metadata.MD, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Header")
-	}
-
-	var r0 metadata.MD
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_IMAMeasurementsClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
-type AgentService_IMAMeasurementsClient_Header_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Header is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Header() *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Header_Call[Res]{Call: _e.mock.On("Header")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Header_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Header_Call[Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Header_Call[Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Recv provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Recv() (*agent.IMAMeasurementsResponse, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Recv")
-	}
-
-	var r0 *agent.IMAMeasurementsResponse
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*agent.IMAMeasurementsResponse, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *agent.IMAMeasurementsResponse); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*agent.IMAMeasurementsResponse)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_IMAMeasurementsClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
-type AgentService_IMAMeasurementsClient_Recv_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Recv is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Recv() *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Recv_Call[Res]{Call: _e.mock.On("Recv")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Recv_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Recv_Call[Res]) Return(_a0 *agent.IMAMeasurementsResponse, _a1 error) *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Recv_Call[Res]) RunAndReturn(run func() (*agent.IMAMeasurementsResponse, error)) *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// RecvMsg provides a mock function with given fields: m
-func (_m *AgentService_IMAMeasurementsClient[Res]) RecvMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for RecvMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
-type AgentService_IMAMeasurementsClient_RecvMsg_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// RecvMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) RecvMsg(m interface{}) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]{Call: _e.mock.On("RecvMsg", m)}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]) Run(run func(m interface{})) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]) Return(_a0 error) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]) RunAndReturn(run func(interface{}) error) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// SendMsg provides a mock function with given fields: m
-func (_m *AgentService_IMAMeasurementsClient[Res]) SendMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
-type AgentService_IMAMeasurementsClient_SendMsg_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// SendMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) SendMsg(m interface{}) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_SendMsg_Call[Res]{Call: _e.mock.On("SendMsg", m)}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call[Res]) Run(run func(m interface{})) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call[Res]) Return(_a0 error) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call[Res]) RunAndReturn(run func(interface{}) error) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Trailer provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Trailer() metadata.MD {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Trailer")
-	}
-
-	var r0 metadata.MD
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
-type AgentService_IMAMeasurementsClient_Trailer_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Trailer is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Trailer() *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Trailer_Call[Res]{Call: _e.mock.On("Trailer")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Trailer_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Trailer_Call[Res]) Return(_a0 metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Trailer_Call[Res]) RunAndReturn(run func() metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewAgentService_IMAMeasurementsClient creates a new instance of AgentService_IMAMeasurementsClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewAgentService_IMAMeasurementsClient[Res interface{}](t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentService_IMAMeasurementsClient[Res] {
-	mock := &AgentService_IMAMeasurementsClient[Res]{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agentservice_algoclient.go

@@ -0,0 +1,442 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package mocks
+
+import (
+	"context"
+
+	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+// NewAgentService_AlgoClient creates a new instance of AgentService_AlgoClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAgentService_AlgoClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentService_AlgoClient {
+	mock := &AgentService_AlgoClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
+type AgentService_AlgoClient struct {
+	mock.Mock
+}
+
+type AgentService_AlgoClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentService_AlgoClient) EXPECT() *AgentService_AlgoClient_Expecter {
+	return &AgentService_AlgoClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseAndRecv provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) CloseAndRecv() (*agent.AlgoResponse, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseAndRecv")
+	}
+
+	var r0 *agent.AlgoResponse
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.AlgoResponse)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_AlgoClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
+type AgentService_AlgoClient_CloseAndRecv_Call struct {
+	*mock.Call
+}
+
+// CloseAndRecv is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) CloseAndRecv() *AgentService_AlgoClient_CloseAndRecv_Call {
+	return &AgentService_AlgoClient_CloseAndRecv_Call{Call: _e.mock.On("CloseAndRecv")}
+}
+
+func (_c *AgentService_AlgoClient_CloseAndRecv_Call) Run(run func()) *AgentService_AlgoClient_CloseAndRecv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseAndRecv_Call) Return(algoResponse *agent.AlgoResponse, err error) *AgentService_AlgoClient_CloseAndRecv_Call {
+	_c.Call.Return(algoResponse, err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseAndRecv_Call) RunAndReturn(run func() (*agent.AlgoResponse, error)) *AgentService_AlgoClient_CloseAndRecv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// CloseSend provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) CloseSend() error {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type AgentService_AlgoClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) CloseSend() *AgentService_AlgoClient_CloseSend_Call {
+	return &AgentService_AlgoClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *AgentService_AlgoClient_CloseSend_Call) Run(run func()) *AgentService_AlgoClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseSend_Call) Return(err error) *AgentService_AlgoClient_CloseSend_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_AlgoClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Context() context.Context {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type AgentService_AlgoClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) Context() *AgentService_AlgoClient_Context_Call {
+	return &AgentService_AlgoClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *AgentService_AlgoClient_Context_Call) Run(run func()) *AgentService_AlgoClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Context_Call) Return(context1 context.Context) *AgentService_AlgoClient_Context_Call {
+	_c.Call.Return(context1)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_AlgoClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Header() (metadata.MD, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_AlgoClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type AgentService_AlgoClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) Header() *AgentService_AlgoClient_Header_Call {
+	return &AgentService_AlgoClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *AgentService_AlgoClient_Header_Call) Run(run func()) *AgentService_AlgoClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_AlgoClient_Header_Call {
+	_c.Call.Return(mD, err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_AlgoClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) RecvMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type AgentService_AlgoClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_AlgoClient_Expecter) RecvMsg(m interface{}) *AgentService_AlgoClient_RecvMsg_Call {
+	return &AgentService_AlgoClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *AgentService_AlgoClient_RecvMsg_Call) Run(run func(m any)) *AgentService_AlgoClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_RecvMsg_Call) Return(err error) *AgentService_AlgoClient_RecvMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_AlgoClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Send provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Send(algoRequest *agent.AlgoRequest) error {
+	ret := _mock.Called(algoRequest)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Send")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
+		r0 = returnFunc(algoRequest)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
+type AgentService_AlgoClient_Send_Call struct {
+	*mock.Call
+}
+
+// Send is a helper method to define mock.On call
+//   - algoRequest *agent.AlgoRequest
+func (_e *AgentService_AlgoClient_Expecter) Send(algoRequest interface{}) *AgentService_AlgoClient_Send_Call {
+	return &AgentService_AlgoClient_Send_Call{Call: _e.mock.On("Send", algoRequest)}
+}
+
+func (_c *AgentService_AlgoClient_Send_Call) Run(run func(algoRequest *agent.AlgoRequest)) *AgentService_AlgoClient_Send_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *agent.AlgoRequest
+		if args[0] != nil {
+			arg0 = args[0].(*agent.AlgoRequest)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Send_Call) Return(err error) *AgentService_AlgoClient_Send_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Send_Call) RunAndReturn(run func(algoRequest *agent.AlgoRequest) error) *AgentService_AlgoClient_Send_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) SendMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type AgentService_AlgoClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_AlgoClient_Expecter) SendMsg(m interface{}) *AgentService_AlgoClient_SendMsg_Call {
+	return &AgentService_AlgoClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *AgentService_AlgoClient_SendMsg_Call) Run(run func(m any)) *AgentService_AlgoClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_SendMsg_Call) Return(err error) *AgentService_AlgoClient_SendMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_AlgoClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Trailer() metadata.MD {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type AgentService_AlgoClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) Trailer() *AgentService_AlgoClient_Trailer_Call {
+	return &AgentService_AlgoClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *AgentService_AlgoClient_Trailer_Call) Run(run func()) *AgentService_AlgoClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Trailer_Call) Return(mD metadata.MD) *AgentService_AlgoClient_Trailer_Call {
+	_c.Call.Return(mD)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_AlgoClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/mocks/agentservice_dataclient.go

@@ -0,0 +1,442 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package mocks
+
+import (
+	"context"
+
+	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+// NewAgentService_DataClient creates a new instance of AgentService_DataClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAgentService_DataClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentService_DataClient {
+	mock := &AgentService_DataClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
+type AgentService_DataClient struct {
+	mock.Mock
+}
+
+type AgentService_DataClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentService_DataClient) EXPECT() *AgentService_DataClient_Expecter {
+	return &AgentService_DataClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseAndRecv provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) CloseAndRecv() (*agent.DataResponse, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseAndRecv")
+	}
+
+	var r0 *agent.DataResponse
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() *agent.DataResponse); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.DataResponse)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_DataClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
+type AgentService_DataClient_CloseAndRecv_Call struct {
+	*mock.Call
+}
+
+// CloseAndRecv is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) CloseAndRecv() *AgentService_DataClient_CloseAndRecv_Call {
+	return &AgentService_DataClient_CloseAndRecv_Call{Call: _e.mock.On("CloseAndRecv")}
+}
+
+func (_c *AgentService_DataClient_CloseAndRecv_Call) Run(run func()) *AgentService_DataClient_CloseAndRecv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseAndRecv_Call) Return(dataResponse *agent.DataResponse, err error) *AgentService_DataClient_CloseAndRecv_Call {
+	_c.Call.Return(dataResponse, err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseAndRecv_Call) RunAndReturn(run func() (*agent.DataResponse, error)) *AgentService_DataClient_CloseAndRecv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// CloseSend provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) CloseSend() error {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type AgentService_DataClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) CloseSend() *AgentService_DataClient_CloseSend_Call {
+	return &AgentService_DataClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *AgentService_DataClient_CloseSend_Call) Run(run func()) *AgentService_DataClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseSend_Call) Return(err error) *AgentService_DataClient_CloseSend_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_DataClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Context() context.Context {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+	return r0
+}
+
+// AgentService_DataClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type AgentService_DataClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) Context() *AgentService_DataClient_Context_Call {
+	return &AgentService_DataClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *AgentService_DataClient_Context_Call) Run(run func()) *AgentService_DataClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Context_Call) Return(context1 context.Context) *AgentService_DataClient_Context_Call {
+	_c.Call.Return(context1)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_DataClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Header() (metadata.MD, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_DataClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type AgentService_DataClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) Header() *AgentService_DataClient_Header_Call {
+	return &AgentService_DataClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *AgentService_DataClient_Header_Call) Run(run func()) *AgentService_DataClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_DataClient_Header_Call {
+	_c.Call.Return(mD, err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_DataClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) RecvMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type AgentService_DataClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_DataClient_Expecter) RecvMsg(m interface{}) *AgentService_DataClient_RecvMsg_Call {
+	return &AgentService_DataClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *AgentService_DataClient_RecvMsg_Call) Run(run func(m any)) *AgentService_DataClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_RecvMsg_Call) Return(err error) *AgentService_DataClient_RecvMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_DataClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Send provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Send(dataRequest *agent.DataRequest) error {
+	ret := _mock.Called(dataRequest)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Send")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
+		r0 = returnFunc(dataRequest)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
+type AgentService_DataClient_Send_Call struct {
+	*mock.Call
+}
+
+// Send is a helper method to define mock.On call
+//   - dataRequest *agent.DataRequest
+func (_e *AgentService_DataClient_Expecter) Send(dataRequest interface{}) *AgentService_DataClient_Send_Call {
+	return &AgentService_DataClient_Send_Call{Call: _e.mock.On("Send", dataRequest)}
+}
+
+func (_c *AgentService_DataClient_Send_Call) Run(run func(dataRequest *agent.DataRequest)) *AgentService_DataClient_Send_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *agent.DataRequest
+		if args[0] != nil {
+			arg0 = args[0].(*agent.DataRequest)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Send_Call) Return(err error) *AgentService_DataClient_Send_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Send_Call) RunAndReturn(run func(dataRequest *agent.DataRequest) error) *AgentService_DataClient_Send_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) SendMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type AgentService_DataClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_DataClient_Expecter) SendMsg(m interface{}) *AgentService_DataClient_SendMsg_Call {
+	return &AgentService_DataClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *AgentService_DataClient_SendMsg_Call) Run(run func(m any)) *AgentService_DataClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_SendMsg_Call) Return(err error) *AgentService_DataClient_SendMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_DataClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Trailer() metadata.MD {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	return r0
+}
+
+// AgentService_DataClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type AgentService_DataClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) Trailer() *AgentService_DataClient_Trailer_Call {
+	return &AgentService_DataClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *AgentService_DataClient_Trailer_Call) Run(run func()) *AgentService_DataClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Trailer_Call) Return(mD metadata.MD) *AgentService_DataClient_Trailer_Call {
+	_c.Call.Return(mD)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_DataClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/mocks/agentservice_imameasurementsclient.go

@@ -0,0 +1,391 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package mocks
+
+import (
+	"context"
+
+	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+// NewAgentService_IMAMeasurementsClient creates a new instance of AgentService_IMAMeasurementsClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAgentService_IMAMeasurementsClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentService_IMAMeasurementsClient {
+	mock := &AgentService_IMAMeasurementsClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// AgentService_IMAMeasurementsClient is an autogenerated mock type for the AgentService_IMAMeasurementsClient type
+type AgentService_IMAMeasurementsClient struct {
+	mock.Mock
+}
+
+type AgentService_IMAMeasurementsClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentService_IMAMeasurementsClient) EXPECT() *AgentService_IMAMeasurementsClient_Expecter {
+	return &AgentService_IMAMeasurementsClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseSend provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) CloseSend() error {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type AgentService_IMAMeasurementsClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) CloseSend() *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	return &AgentService_IMAMeasurementsClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) Run(run func()) *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) Return(err error) *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Context() context.Context {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type AgentService_IMAMeasurementsClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Context() *AgentService_IMAMeasurementsClient_Context_Call {
+	return &AgentService_IMAMeasurementsClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Context_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Context_Call) Return(context1 context.Context) *AgentService_IMAMeasurementsClient_Context_Call {
+	_c.Call.Return(context1)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_IMAMeasurementsClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Header() (metadata.MD, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_IMAMeasurementsClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type AgentService_IMAMeasurementsClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Header() *AgentService_IMAMeasurementsClient_Header_Call {
+	return &AgentService_IMAMeasurementsClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Header_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_IMAMeasurementsClient_Header_Call {
+	_c.Call.Return(mD, err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_IMAMeasurementsClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Recv provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Recv() (*agent.IMAMeasurementsResponse, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Recv")
+	}
+
+	var r0 *agent.IMAMeasurementsResponse
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (*agent.IMAMeasurementsResponse, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() *agent.IMAMeasurementsResponse); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.IMAMeasurementsResponse)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_IMAMeasurementsClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
+type AgentService_IMAMeasurementsClient_Recv_Call struct {
+	*mock.Call
+}
+
+// Recv is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Recv() *AgentService_IMAMeasurementsClient_Recv_Call {
+	return &AgentService_IMAMeasurementsClient_Recv_Call{Call: _e.mock.On("Recv")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Recv_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Recv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Recv_Call) Return(iMAMeasurementsResponse *agent.IMAMeasurementsResponse, err error) *AgentService_IMAMeasurementsClient_Recv_Call {
+	_c.Call.Return(iMAMeasurementsResponse, err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Recv_Call) RunAndReturn(run func() (*agent.IMAMeasurementsResponse, error)) *AgentService_IMAMeasurementsClient_Recv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) RecvMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type AgentService_IMAMeasurementsClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_IMAMeasurementsClient_Expecter) RecvMsg(m interface{}) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	return &AgentService_IMAMeasurementsClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) Run(run func(m any)) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) Return(err error) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) SendMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type AgentService_IMAMeasurementsClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_IMAMeasurementsClient_Expecter) SendMsg(m interface{}) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	return &AgentService_IMAMeasurementsClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) Run(run func(m any)) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) Return(err error) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Trailer() metadata.MD {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type AgentService_IMAMeasurementsClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Trailer() *AgentService_IMAMeasurementsClient_Trailer_Call {
+	return &AgentService_IMAMeasurementsClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) Return(mD metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call {
+	_c.Call.Return(mD)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/mocks/service.go

@@ -1,19 +1,34 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	agent "github.com/ultravioletrs/cocos/agent"
-	attestation "github.com/ultravioletrs/cocos/pkg/attestation"
-
-	context "context"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
 )
 
+// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewService(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Service is an autogenerated mock type for the Service type
 type Service struct {
 	mock.Mock
@@ -27,21 +42,20 @@ func (_m *Service) EXPECT() *Service_Expecter {
 	return &Service_Expecter{mock: &_m.Mock}
 }
 
-// Algo provides a mock function with given fields: ctx, algorithm
-func (_m *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
-	ret := _m.Called(ctx, algorithm)
+// Algo provides a mock function for the type Service
+func (_mock *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
+	ret := _mock.Called(ctx, algorithm)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Algo")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
-		r0 = rf(ctx, algorithm)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
+		r0 = returnFunc(ctx, algorithm)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -59,24 +73,35 @@ func (_e *Service_Expecter) Algo(ctx interface{}, algorithm interface{}) *Servic
 
 func (_c *Service_Algo_Call) Run(run func(ctx context.Context, algorithm agent.Algorithm)) *Service_Algo_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(agent.Algorithm))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 agent.Algorithm
+		if args[1] != nil {
+			arg1 = args[1].(agent.Algorithm)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Algo_Call) Return(_a0 error) *Service_Algo_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_Algo_Call) Return(err error) *Service_Algo_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_Algo_Call) RunAndReturn(run func(context.Context, agent.Algorithm) error) *Service_Algo_Call {
+func (_c *Service_Algo_Call) RunAndReturn(run func(ctx context.Context, algorithm agent.Algorithm) error) *Service_Algo_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Attestation provides a mock function with given fields: ctx, reportData, nonce, attType
-func (_m *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
-	ret := _m.Called(ctx, reportData, nonce, attType)
+// Attestation provides a mock function for the type Service
+func (_mock *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
+	ret := _mock.Called(ctx, reportData, nonce, attType)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Attestation")
@@ -84,23 +109,21 @@ func (_m *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
-		return rf(ctx, reportData, nonce, attType)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
+		return returnFunc(ctx, reportData, nonce, attType)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) []byte); ok {
-		r0 = rf(ctx, reportData, nonce, attType)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) []byte); ok {
+		r0 = returnFunc(ctx, reportData, nonce, attType)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) error); ok {
-		r1 = rf(ctx, reportData, nonce, attType)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) error); ok {
+		r1 = returnFunc(ctx, reportData, nonce, attType)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -120,24 +143,45 @@ func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{},
 
 func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType)) *Service_Attestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].([64]byte), args[2].([32]byte), args[3].(attestation.PlatformType))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 [64]byte
+		if args[1] != nil {
+			arg1 = args[1].([64]byte)
+		}
+		var arg2 [32]byte
+		if args[2] != nil {
+			arg2 = args[2].([32]byte)
+		}
+		var arg3 attestation.PlatformType
+		if args[3] != nil {
+			arg3 = args[3].(attestation.PlatformType)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+			arg3,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Attestation_Call) Return(_a0 []byte, _a1 error) *Service_Attestation_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_Attestation_Call) Return(bytes []byte, err error) *Service_Attestation_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)) *Service_Attestation_Call {
+func (_c *Service_Attestation_Call) RunAndReturn(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error)) *Service_Attestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// AzureAttestationToken provides a mock function with given fields: ctx, nonce
-func (_m *Service) AzureAttestationToken(ctx context.Context, nonce [32]byte) ([]byte, error) {
-	ret := _m.Called(ctx, nonce)
+// AzureAttestationToken provides a mock function for the type Service
+func (_mock *Service) AzureAttestationToken(ctx context.Context, nonce [32]byte) ([]byte, error) {
+	ret := _mock.Called(ctx, nonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AzureAttestationToken")
@@ -145,23 +189,21 @@ func (_m *Service) AzureAttestationToken(ctx context.Context, nonce [32]byte) ([
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, [32]byte) ([]byte, error)); ok {
-		return rf(ctx, nonce)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) ([]byte, error)); ok {
+		return returnFunc(ctx, nonce)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, [32]byte) []byte); ok {
-		r0 = rf(ctx, nonce)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) []byte); ok {
+		r0 = returnFunc(ctx, nonce)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, [32]byte) error); ok {
-		r1 = rf(ctx, nonce)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, [32]byte) error); ok {
+		r1 = returnFunc(ctx, nonce)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -179,36 +221,46 @@ func (_e *Service_Expecter) AzureAttestationToken(ctx interface{}, nonce interfa
 
 func (_c *Service_AzureAttestationToken_Call) Run(run func(ctx context.Context, nonce [32]byte)) *Service_AzureAttestationToken_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].([32]byte))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 [32]byte
+		if args[1] != nil {
+			arg1 = args[1].([32]byte)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_AzureAttestationToken_Call) Return(_a0 []byte, _a1 error) *Service_AzureAttestationToken_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_AzureAttestationToken_Call) Return(bytes []byte, err error) *Service_AzureAttestationToken_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_AzureAttestationToken_Call) RunAndReturn(run func(context.Context, [32]byte) ([]byte, error)) *Service_AzureAttestationToken_Call {
+func (_c *Service_AzureAttestationToken_Call) RunAndReturn(run func(ctx context.Context, nonce [32]byte) ([]byte, error)) *Service_AzureAttestationToken_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Data provides a mock function with given fields: ctx, dataset
-func (_m *Service) Data(ctx context.Context, dataset agent.Dataset) error {
-	ret := _m.Called(ctx, dataset)
+// Data provides a mock function for the type Service
+func (_mock *Service) Data(ctx context.Context, dataset agent.Dataset) error {
+	ret := _mock.Called(ctx, dataset)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Data")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
-		r0 = rf(ctx, dataset)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
+		r0 = returnFunc(ctx, dataset)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -226,24 +278,35 @@ func (_e *Service_Expecter) Data(ctx interface{}, dataset interface{}) *Service_
 
 func (_c *Service_Data_Call) Run(run func(ctx context.Context, dataset agent.Dataset)) *Service_Data_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(agent.Dataset))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 agent.Dataset
+		if args[1] != nil {
+			arg1 = args[1].(agent.Dataset)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Data_Call) Return(_a0 error) *Service_Data_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_Data_Call) Return(err error) *Service_Data_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_Data_Call) RunAndReturn(run func(context.Context, agent.Dataset) error) *Service_Data_Call {
+func (_c *Service_Data_Call) RunAndReturn(run func(ctx context.Context, dataset agent.Dataset) error) *Service_Data_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// IMAMeasurements provides a mock function with given fields: ctx
-func (_m *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
-	ret := _m.Called(ctx)
+// IMAMeasurements provides a mock function for the type Service
+func (_mock *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for IMAMeasurements")
@@ -252,31 +315,28 @@ func (_m *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error)
 	var r0 []byte
 	var r1 []byte
 	var r2 error
-	if rf, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
-		return rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
+		return returnFunc(ctx)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) []byte); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context) []byte); ok {
-		r1 = rf(ctx)
+	if returnFunc, ok := ret.Get(1).(func(context.Context) []byte); ok {
+		r1 = returnFunc(ctx)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(2).(func(context.Context) error); ok {
-		r2 = rf(ctx)
+	if returnFunc, ok := ret.Get(2).(func(context.Context) error); ok {
+		r2 = returnFunc(ctx)
 	} else {
 		r2 = ret.Error(2)
 	}
-
 	return r0, r1, r2
 }
 
@@ -293,36 +353,41 @@ func (_e *Service_Expecter) IMAMeasurements(ctx interface{}) *Service_IMAMeasure
 
 func (_c *Service_IMAMeasurements_Call) Run(run func(ctx context.Context)) *Service_IMAMeasurements_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_IMAMeasurements_Call) Return(_a0 []byte, _a1 []byte, _a2 error) *Service_IMAMeasurements_Call {
-	_c.Call.Return(_a0, _a1, _a2)
+func (_c *Service_IMAMeasurements_Call) Return(bytes []byte, bytes1 []byte, err error) *Service_IMAMeasurements_Call {
+	_c.Call.Return(bytes, bytes1, err)
 	return _c
 }
 
-func (_c *Service_IMAMeasurements_Call) RunAndReturn(run func(context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
+func (_c *Service_IMAMeasurements_Call) RunAndReturn(run func(ctx context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// InitComputation provides a mock function with given fields: ctx, cmp
-func (_m *Service) InitComputation(ctx context.Context, cmp agent.Computation) error {
-	ret := _m.Called(ctx, cmp)
+// InitComputation provides a mock function for the type Service
+func (_mock *Service) InitComputation(ctx context.Context, cmp agent.Computation) error {
+	ret := _mock.Called(ctx, cmp)
 
 	if len(ret) == 0 {
 		panic("no return value specified for InitComputation")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
-		r0 = rf(ctx, cmp)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
+		r0 = returnFunc(ctx, cmp)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -340,24 +405,35 @@ func (_e *Service_Expecter) InitComputation(ctx interface{}, cmp interface{}) *S
 
 func (_c *Service_InitComputation_Call) Run(run func(ctx context.Context, cmp agent.Computation)) *Service_InitComputation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(agent.Computation))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 agent.Computation
+		if args[1] != nil {
+			arg1 = args[1].(agent.Computation)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_InitComputation_Call) Return(_a0 error) *Service_InitComputation_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_InitComputation_Call) Return(err error) *Service_InitComputation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_InitComputation_Call) RunAndReturn(run func(context.Context, agent.Computation) error) *Service_InitComputation_Call {
+func (_c *Service_InitComputation_Call) RunAndReturn(run func(ctx context.Context, cmp agent.Computation) error) *Service_InitComputation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Result provides a mock function with given fields: ctx
-func (_m *Service) Result(ctx context.Context) ([]byte, error) {
-	ret := _m.Called(ctx)
+// Result provides a mock function for the type Service
+func (_mock *Service) Result(ctx context.Context) ([]byte, error) {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Result")
@@ -365,23 +441,21 @@ func (_m *Service) Result(ctx context.Context) ([]byte, error) {
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
-		return rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
+		return returnFunc(ctx)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) []byte); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
-		r1 = rf(ctx)
+	if returnFunc, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = returnFunc(ctx)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -398,36 +472,41 @@ func (_e *Service_Expecter) Result(ctx interface{}) *Service_Result_Call {
 
 func (_c *Service_Result_Call) Run(run func(ctx context.Context)) *Service_Result_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Result_Call) Return(_a0 []byte, _a1 error) *Service_Result_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_Result_Call) Return(bytes []byte, err error) *Service_Result_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_Result_Call) RunAndReturn(run func(context.Context) ([]byte, error)) *Service_Result_Call {
+func (_c *Service_Result_Call) RunAndReturn(run func(ctx context.Context) ([]byte, error)) *Service_Result_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// State provides a mock function with no fields
-func (_m *Service) State() string {
-	ret := _m.Called()
+// State provides a mock function for the type Service
+func (_mock *Service) State() string {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for State")
 	}
 
 	var r0 string
-	if rf, ok := ret.Get(0).(func() string); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() string); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Get(0).(string)
 	}
-
 	return r0
 }
 
@@ -448,8 +527,8 @@ func (_c *Service_State_Call) Run(run func()) *Service_State_Call {
 	return _c
 }
 
-func (_c *Service_State_Call) Return(_a0 string) *Service_State_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_State_Call) Return(s string) *Service_State_Call {
+	_c.Call.Return(s)
 	return _c
 }
 
@@ -458,21 +537,20 @@ func (_c *Service_State_Call) RunAndReturn(run func() string) *Service_State_Cal
 	return _c
 }
 
-// StopComputation provides a mock function with given fields: ctx
-func (_m *Service) StopComputation(ctx context.Context) error {
-	ret := _m.Called(ctx)
+// StopComputation provides a mock function for the type Service
+func (_mock *Service) StopComputation(ctx context.Context) error {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for StopComputation")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context) error); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -489,31 +567,23 @@ func (_e *Service_Expecter) StopComputation(ctx interface{}) *Service_StopComput
 
 func (_c *Service_StopComputation_Call) Run(run func(ctx context.Context)) *Service_StopComputation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_StopComputation_Call) Return(_a0 error) *Service_StopComputation_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_StopComputation_Call) Return(err error) *Service_StopComputation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_StopComputation_Call) RunAndReturn(run func(context.Context) error) *Service_StopComputation_Call {
+func (_c *Service_StopComputation_Call) RunAndReturn(run func(ctx context.Context) error) *Service_StopComputation_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewService(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Service {
-	mock := &Service{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/statemachine/mocks/statemachine.go

@@ -1,17 +1,33 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	context "context"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
-	statemachine "github.com/ultravioletrs/cocos/agent/statemachine"
+	"github.com/ultravioletrs/cocos/agent/statemachine"
 )
 
+// NewStateMachine creates a new instance of StateMachine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewStateMachine(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *StateMachine {
+	mock := &StateMachine{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // StateMachine is an autogenerated mock type for the StateMachine type
 type StateMachine struct {
 	mock.Mock
@@ -25,9 +41,10 @@ func (_m *StateMachine) EXPECT() *StateMachine_Expecter {
 	return &StateMachine_Expecter{mock: &_m.Mock}
 }
 
-// AddTransition provides a mock function with given fields: t
-func (_m *StateMachine) AddTransition(t statemachine.Transition) {
-	_m.Called(t)
+// AddTransition provides a mock function for the type StateMachine
+func (_mock *StateMachine) AddTransition(t statemachine.Transition) {
+	_mock.Called(t)
+	return
 }
 
 // StateMachine_AddTransition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddTransition'
@@ -43,7 +60,13 @@ func (_e *StateMachine_Expecter) AddTransition(t interface{}) *StateMachine_AddT
 
 func (_c *StateMachine_AddTransition_Call) Run(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.Transition))
+		var arg0 statemachine.Transition
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.Transition)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
@@ -53,28 +76,27 @@ func (_c *StateMachine_AddTransition_Call) Return() *StateMachine_AddTransition_
 	return _c
 }
 
-func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(statemachine.Transition)) *StateMachine_AddTransition_Call {
+func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
 	_c.Run(run)
 	return _c
 }
 
-// GetState provides a mock function with no fields
-func (_m *StateMachine) GetState() statemachine.State {
-	ret := _m.Called()
+// GetState provides a mock function for the type StateMachine
+func (_mock *StateMachine) GetState() statemachine.State {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for GetState")
 	}
 
 	var r0 statemachine.State
-	if rf, ok := ret.Get(0).(func() statemachine.State); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() statemachine.State); ok {
+		r0 = returnFunc()
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(statemachine.State)
 		}
 	}
-
 	return r0
 }
 
@@ -95,8 +117,8 @@ func (_c *StateMachine_GetState_Call) Run(run func()) *StateMachine_GetState_Cal
 	return _c
 }
 
-func (_c *StateMachine_GetState_Call) Return(_a0 statemachine.State) *StateMachine_GetState_Call {
-	_c.Call.Return(_a0)
+func (_c *StateMachine_GetState_Call) Return(state statemachine.State) *StateMachine_GetState_Call {
+	_c.Call.Return(state)
 	return _c
 }
 
@@ -105,9 +127,10 @@ func (_c *StateMachine_GetState_Call) RunAndReturn(run func() statemachine.State
 	return _c
 }
 
-// Reset provides a mock function with given fields: initialState
-func (_m *StateMachine) Reset(initialState statemachine.State) {
-	_m.Called(initialState)
+// Reset provides a mock function for the type StateMachine
+func (_mock *StateMachine) Reset(initialState statemachine.State) {
+	_mock.Called(initialState)
+	return
 }
 
 // StateMachine_Reset_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Reset'
@@ -123,7 +146,13 @@ func (_e *StateMachine_Expecter) Reset(initialState interface{}) *StateMachine_R
 
 func (_c *StateMachine_Reset_Call) Run(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.State))
+		var arg0 statemachine.State
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.State)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
@@ -133,14 +162,15 @@ func (_c *StateMachine_Reset_Call) Return() *StateMachine_Reset_Call {
 	return _c
 }
 
-func (_c *StateMachine_Reset_Call) RunAndReturn(run func(statemachine.State)) *StateMachine_Reset_Call {
+func (_c *StateMachine_Reset_Call) RunAndReturn(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
 	_c.Run(run)
 	return _c
 }
 
-// SendEvent provides a mock function with given fields: event
-func (_m *StateMachine) SendEvent(event statemachine.Event) {
-	_m.Called(event)
+// SendEvent provides a mock function for the type StateMachine
+func (_mock *StateMachine) SendEvent(event statemachine.Event) {
+	_mock.Called(event)
+	return
 }
 
 // StateMachine_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
@@ -156,7 +186,13 @@ func (_e *StateMachine_Expecter) SendEvent(event interface{}) *StateMachine_Send
 
 func (_c *StateMachine_SendEvent_Call) Run(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.Event))
+		var arg0 statemachine.Event
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.Event)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
@@ -166,14 +202,15 @@ func (_c *StateMachine_SendEvent_Call) Return() *StateMachine_SendEvent_Call {
 	return _c
 }
 
-func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(statemachine.Event)) *StateMachine_SendEvent_Call {
+func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
 	_c.Run(run)
 	return _c
 }
 
-// SetAction provides a mock function with given fields: state, action
-func (_m *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
-	_m.Called(state, action)
+// SetAction provides a mock function for the type StateMachine
+func (_mock *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
+	_mock.Called(state, action)
+	return
 }
 
 // StateMachine_SetAction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetAction'
@@ -190,7 +227,18 @@ func (_e *StateMachine_Expecter) SetAction(state interface{}, action interface{}
 
 func (_c *StateMachine_SetAction_Call) Run(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.State), args[1].(statemachine.Action))
+		var arg0 statemachine.State
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.State)
+		}
+		var arg1 statemachine.Action
+		if args[1] != nil {
+			arg1 = args[1].(statemachine.Action)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
@@ -200,26 +248,25 @@ func (_c *StateMachine_SetAction_Call) Return() *StateMachine_SetAction_Call {
 	return _c
 }
 
-func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(statemachine.State, statemachine.Action)) *StateMachine_SetAction_Call {
+func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
 	_c.Run(run)
 	return _c
 }
 
-// Start provides a mock function with given fields: ctx
-func (_m *StateMachine) Start(ctx context.Context) error {
-	ret := _m.Called(ctx)
+// Start provides a mock function for the type StateMachine
+func (_mock *StateMachine) Start(ctx context.Context) error {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Start")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context) error); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -236,31 +283,23 @@ func (_e *StateMachine_Expecter) Start(ctx interface{}) *StateMachine_Start_Call
 
 func (_c *StateMachine_Start_Call) Run(run func(ctx context.Context)) *StateMachine_Start_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *StateMachine_Start_Call) Return(_a0 error) *StateMachine_Start_Call {
-	_c.Call.Return(_a0)
+func (_c *StateMachine_Start_Call) Return(err error) *StateMachine_Start_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *StateMachine_Start_Call) RunAndReturn(run func(context.Context) error) *StateMachine_Start_Call {
+func (_c *StateMachine_Start_Call) RunAndReturn(run func(ctx context.Context) error) *StateMachine_Start_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewStateMachine creates a new instance of StateMachine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
-// The first argument is typically a *testing.T value.
-func NewStateMachine(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *StateMachine {
-	mock := &StateMachine{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}