1 /*
2  * Copyright (c) Facebook, Inc. and its affiliates.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <folly/Memory.h>
20 #include <folly/io/async/AsyncTransport.h>
21 #include <folly/portability/GMock.h>
22 
23 namespace folly {
24 namespace test {
25 
26 class MockAsyncTransport : public AsyncTransport {
27  public:
28   MOCK_METHOD1(setEventCallback, void(EventRecvmsgCallback*));
29   MOCK_METHOD1(setReadCB, void(ReadCallback*));
30   MOCK_CONST_METHOD0(getReadCallback, ReadCallback*());
31   MOCK_CONST_METHOD0(getReadCB, ReadCallback*());
32   MOCK_METHOD4(write, void(WriteCallback*, const void*, size_t, WriteFlags));
33   MOCK_METHOD4(writev, void(WriteCallback*, const iovec*, size_t, WriteFlags));
34   MOCK_METHOD3(
35       writeChain,
36       void(WriteCallback*, std::shared_ptr<folly::IOBuf>, WriteFlags));
37 
38   void writeChain(
39       WriteCallback* callback,
40       std::unique_ptr<folly::IOBuf>&& iob,
41       WriteFlags flags = WriteFlags::NONE) override {
42     writeChain(callback, std::shared_ptr<folly::IOBuf>(iob.release()), flags);
43   }
44 
45   MOCK_METHOD0(close, void());
46   MOCK_METHOD0(closeNow, void());
47   MOCK_METHOD0(closeWithReset, void());
48   MOCK_METHOD0(shutdownWrite, void());
49   MOCK_METHOD0(shutdownWriteNow, void());
50   MOCK_CONST_METHOD0(good, bool());
51   MOCK_CONST_METHOD0(readable, bool());
52   MOCK_CONST_METHOD0(connecting, bool());
53   MOCK_CONST_METHOD0(error, bool());
54   MOCK_METHOD1(attachEventBase, void(EventBase*));
55   MOCK_METHOD0(detachEventBase, void());
56   MOCK_CONST_METHOD0(isDetachable, bool());
57   MOCK_CONST_METHOD0(getEventBase, EventBase*());
58   MOCK_METHOD1(setSendTimeout, void(uint32_t));
59   MOCK_CONST_METHOD0(getSendTimeout, uint32_t());
60   MOCK_CONST_METHOD1(getLocalAddress, void(folly::SocketAddress*));
61   MOCK_CONST_METHOD1(getPeerAddress, void(folly::SocketAddress*));
62   MOCK_CONST_METHOD0(getAppBytesWritten, size_t());
63   MOCK_CONST_METHOD0(getRawBytesWritten, size_t());
64   MOCK_CONST_METHOD0(getAppBytesReceived, size_t());
65   MOCK_CONST_METHOD0(getRawBytesReceived, size_t());
66   MOCK_CONST_METHOD0(getAppBytesBuffered, size_t());
67   MOCK_CONST_METHOD0(getRawBytesBuffered, size_t());
68   MOCK_CONST_METHOD0(isEorTrackingEnabled, bool());
69   MOCK_METHOD1(setEorTracking, void(bool));
70   MOCK_CONST_METHOD0(getWrappedTransport, AsyncTransport*());
71   MOCK_CONST_METHOD0(isReplaySafe, bool());
72   MOCK_METHOD1(
73       setReplaySafetyCallback, void(AsyncTransport::ReplaySafetyCallback*));
74   MOCK_CONST_METHOD0(getSecurityProtocol, std::string());
75   MOCK_CONST_METHOD0(getPeerCertificate, const AsyncTransportCertificate*());
76 };
77 
78 class MockReplaySafetyCallback : public AsyncTransport::ReplaySafetyCallback {
79  public:
80   MOCK_METHOD0(onReplaySafe_, void());
onReplaySafe()81   void onReplaySafe() noexcept override { onReplaySafe_(); }
82 };
83 
84 class MockReadCallback : public AsyncTransport::ReadCallback {
85  public:
86   MOCK_METHOD2(getReadBuffer, void(void**, size_t*));
87 
88   MOCK_METHOD1(readDataAvailable_, void(size_t));
readDataAvailable(size_t size)89   void readDataAvailable(size_t size) noexcept override {
90     readDataAvailable_(size);
91   }
92 
93   MOCK_METHOD0(isBufferMovable_, bool());
isBufferMovable()94   bool isBufferMovable() noexcept override { return isBufferMovable_(); }
95 
96   MOCK_METHOD1(readBufferAvailable_, void(std::unique_ptr<folly::IOBuf>&));
readBufferAvailable(std::unique_ptr<folly::IOBuf> readBuf)97   void readBufferAvailable(
98       std::unique_ptr<folly::IOBuf> readBuf) noexcept override {
99     readBufferAvailable_(readBuf);
100   }
101 
102   MOCK_METHOD0(readEOF_, void());
readEOF()103   void readEOF() noexcept override { readEOF_(); }
104 
105   MOCK_METHOD1(readErr_, void(const AsyncSocketException&));
readErr(const AsyncSocketException & ex)106   void readErr(const AsyncSocketException& ex) noexcept override {
107     readErr_(ex);
108   }
109 };
110 
111 class MockWriteCallback : public AsyncTransport::WriteCallback {
112  public:
113   MOCK_METHOD0(writeSuccess_, void());
writeSuccess()114   void writeSuccess() noexcept override { writeSuccess_(); }
115 
116   MOCK_METHOD2(writeErr_, void(size_t, const AsyncSocketException&));
writeErr(size_t size,const AsyncSocketException & ex)117   void writeErr(size_t size, const AsyncSocketException& ex) noexcept override {
118     writeErr_(size, ex);
119   }
120 };
121 
122 } // namespace test
123 } // namespace folly
124